diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py index 6d026fb86..d9f96838a 100644 --- a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py @@ -43,8 +43,8 @@ def _estimate_prompt_len( """ try: from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig - from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( - Qwen3TTSTalkerForConditionalGeneration, + from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + Qwen3TTSPromptEmbedsBuilder, ) if model_name not in _cache: @@ -124,7 +124,7 @@ def _estimate_prompt_len( except Exception: return None - return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( additional_information=additional_information, task_type=task_type, tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"], diff --git a/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py b/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py index 1f71cf8fb..d82e1a68b 100644 --- a/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py +++ b/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py @@ -28,7 +28,7 @@ MODEL_BASE = "Qwen/Qwen3-TTS-12Hz-0.6B-Base" MODEL_BASE_1_7B = "Qwen/Qwen3-TTS-12Hz-1.7B-Base" # A synthetic 1024-dim speaker embedding (all 0.1 — not a real voice, but -# exercises the full code path through the talker's _build_prompt_embeds). +# exercises the full code path through Qwen3TTSPromptEmbedsBuilder.build_prompt_embeds). DUMMY_EMBEDDING_1024 = [0.1] * 1024 DUMMY_EMBEDDING_2048 = [0.1] * 2048 diff --git a/tests/examples/offline_inference/test_qwen3_tts_estimator.py b/tests/examples/offline_inference/test_qwen3_tts_estimator.py index fde638ba5..4f94eef0d 100644 --- a/tests/examples/offline_inference/test_qwen3_tts_estimator.py +++ b/tests/examples/offline_inference/test_qwen3_tts_estimator.py @@ -57,11 +57,11 @@ def wav_file(tmp_path: Path) -> Path: def estimator_cache(monkeypatch): """ Bypass model loading by pre-populating the estimator cache, and shortcut - Qwen3TTSTalker.estimate_prompt_len_from_additional_information so that - the test directly observes the real _estimate_ref_code_len closure. + Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information + so that the test directly observes the real _estimate_ref_code_len closure. """ - from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( - Qwen3TTSTalkerForConditionalGeneration, + from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + Qwen3TTSPromptEmbedsBuilder, ) def _shortcut(**kwargs): @@ -69,7 +69,7 @@ def estimator_cache(monkeypatch): return kwargs["estimate_ref_code_len"](ref_audio) monkeypatch.setattr( - Qwen3TTSTalkerForConditionalGeneration, + Qwen3TTSPromptEmbedsBuilder, "estimate_prompt_len_from_additional_information", staticmethod(_shortcut), ) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 7fe7fc90a..789ed229a 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -2147,8 +2147,8 @@ class OmniRunner: _cache = self._prompt_len_estimate_cache try: from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig - from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( - Qwen3TTSTalkerForConditionalGeneration, + from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + Qwen3TTSPromptEmbedsBuilder, ) if model_name not in _cache: @@ -2160,7 +2160,7 @@ class OmniRunner: tok, tcfg = _cache[model_name] task_type = (additional_information.get("task_type") or ["CustomVoice"])[0] - return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( additional_information=additional_information, task_type=task_type, tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"], diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py index cd590161f..7da61fb33 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py @@ -4,34 +4,105 @@ from types import SimpleNamespace import numpy as np import torch +from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + NORMALIZED_REF_AUDIO_KEY, + PRECOMPUTED_REF_CODE_KEY, + PRECOMPUTED_REF_IDS_KEY, + PRECOMPUTED_TEXT_IDS_KEY, + Qwen3TTSPromptEmbedsBuilder, +) from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( - _NORMALIZED_REF_AUDIO_KEY, - _PRECOMPUTED_REF_CODE_KEY, - _PRECOMPUTED_REF_IDS_KEY, - _PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSTalkerForConditionalGeneration, ) -def _make_minimal_talker(): +def _make_minimal_talker( + tts_pad_embed: torch.Tensor | None = None, + *, + build_prompt_embeds=None, +): + """Construct a bare ``Qwen3TTSTalkerForConditionalGeneration`` for preprocess tests. + + Skips ``__init__`` and ``load_weights`` (so we don't need real config / + checkpoint files) and pre-populates the minimum set of attributes the + preprocess / decode paths read: the persistent ``tts_pad`` buffer and a + stub :class:`Qwen3TTSPromptEmbedsBuilder` with a configurable + ``build_prompt_embeds`` (defaults to raising, since most preprocess + tests don't exercise the builder). + """ model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16) - model._ref_audio_artifact_cache_max_entries = 256 - model._ref_audio_artifact_cache = OrderedDict() + if tts_pad_embed is None: + tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16) + model._tts_pad_embed = tts_pad_embed + + def _default_raise(**_kwargs): + raise AssertionError("build_prompt_embeds was not stubbed in this test") + + model._prompt_builder = SimpleNamespace( + build_prompt_embeds=build_prompt_embeds if build_prompt_embeds is not None else _default_raise, + ) return model +def _make_minimal_builder( + *, + config: SimpleNamespace | None = None, + talker_config: SimpleNamespace | None = None, + tts_pad_embed: torch.Tensor | None = None, +): + """Construct a bare :class:`Qwen3TTSPromptEmbedsBuilder` instance. + + Bypasses ``__init__`` so we don't need a real model checkpoint; tests + populate the attributes they actually exercise. + """ + builder = Qwen3TTSPromptEmbedsBuilder.__new__(Qwen3TTSPromptEmbedsBuilder) + builder._config = config or SimpleNamespace( + tts_bos_token_id=100, + tts_eos_token_id=101, + tts_pad_token_id=102, + ) + builder._talker_config = talker_config or SimpleNamespace( + codec_nothink_id=10, + codec_think_bos_id=11, + codec_think_eos_id=12, + codec_think_id=13, + codec_language_id={}, + codec_pad_id=7, + codec_bos_id=8, + num_code_groups=2, + spk_is_dialect={}, + ) + builder._model_path = "" + builder._text_embedding = None + builder._text_projection = lambda x: x + builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device) + builder._residual_code_embeddings = lambda: [] + builder._speaker_encoder = None + builder._tts_pad_embed_buffer = ( + tts_pad_embed if tts_pad_embed is not None else torch.zeros((1, 4), dtype=torch.bfloat16) + ) + builder._encode_ref_audio_batch_fn = lambda wavs, sr, *, device: (_ for _ in ()).throw( + AssertionError("encode_ref_audio_batch_fn not stubbed") + ) + builder._speaker_cache = None + builder._text_tokenizer = None + builder._ref_audio_artifact_cache_max_entries = 256 + builder._ref_audio_artifact_cache = OrderedDict() + builder._resampler_cache = OrderedDict() + builder._resampler_cache_max = 16 + return builder + + def test_single_token_prefill_uses_prefill_path(): - model = _make_minimal_talker() full_prompt_embeds = torch.arange(12, dtype=torch.float32).reshape(3, 4) trailing_text = torch.ones((2, 4), dtype=torch.float32) - tts_pad = torch.full((1, 4), 0.5, dtype=torch.float32) ref_code = torch.arange(32, dtype=torch.long).reshape(2, 16) def fake_build_prompt_embeds(*, task_type, info_dict): - return full_prompt_embeds, trailing_text, tts_pad, 2, ref_code + return full_prompt_embeds, trailing_text, 2, ref_code - model._build_prompt_embeds = fake_build_prompt_embeds + model = _make_minimal_talker(build_prompt_embeds=fake_build_prompt_embeds) input_ids = torch.tensor([123], dtype=torch.long) out_ids, out_embeds, update = model.preprocess( @@ -50,22 +121,22 @@ def test_single_token_prefill_uses_prefill_path(): assert update["meta"]["talker_text_offset"] == 0 assert update["meta"]["ref_code_len"] == 2 assert torch.equal(update["embed"]["prefill"], full_prompt_embeds) - assert torch.equal(update["embed"]["tts_pad"], tts_pad) + # ``tts_pad`` is no longer round-tripped through ``info_dict`` — it is + # precomputed once into ``model._tts_pad_embed`` at load time. + assert "tts_pad" not in update["embed"] assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text) assert torch.equal(update["codes"]["ref"], ref_code) assert update["codes"]["audio"].shape == (1, 16) def test_single_token_prefill_can_be_inferred_from_token_progress(): - model = _make_minimal_talker() full_prompt_embeds = torch.arange(8, dtype=torch.float32).reshape(2, 4) trailing_text = torch.ones((1, 4), dtype=torch.float32) - tts_pad = torch.zeros((1, 4), dtype=torch.float32) def fake_build_prompt_embeds(*, task_type, info_dict): - return full_prompt_embeds, trailing_text, tts_pad, None, None + return full_prompt_embeds, trailing_text, None, None - model._build_prompt_embeds = fake_build_prompt_embeds + model = _make_minimal_talker(build_prompt_embeds=fake_build_prompt_embeds) out_ids, out_embeds, update = model.preprocess( input_ids=torch.tensor([123], dtype=torch.long), @@ -82,7 +153,8 @@ def test_single_token_prefill_can_be_inferred_from_token_progress(): def test_decode_advances_trailing_text_by_offset_without_rewriting_tail(): - model = _make_minimal_talker() + tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16) + model = _make_minimal_talker(tts_pad_embed=tts_pad) def fake_embed_input_ids(input_ids): return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4) @@ -90,7 +162,6 @@ def test_decode_advances_trailing_text_by_offset_without_rewriting_tail(): model.embed_input_ids = fake_embed_input_ids trailing_text = torch.arange(12, dtype=torch.float32).reshape(3, 4) last_hidden = torch.full((4,), 2.0, dtype=torch.float32) - tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32) out_ids, out_embeds, update = model.preprocess( input_ids=torch.tensor([123], dtype=torch.long), @@ -98,7 +169,6 @@ def test_decode_advances_trailing_text_by_offset_without_rewriting_tail(): text=["hello"], task_type=["CustomVoice"], hidden_states={"trailing_text": trailing_text, "last": last_hidden}, - embed={"tts_pad": tts_pad}, meta={"talker_text_offset": 1}, _omni_is_prefill=False, _omni_num_computed_tokens=2, @@ -115,7 +185,8 @@ def test_decode_advances_trailing_text_by_offset_without_rewriting_tail(): def test_decode_advances_trailing_text_offset_across_multiple_steps(): - model = _make_minimal_talker() + tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16) + model = _make_minimal_talker(tts_pad_embed=tts_pad) def fake_embed_input_ids(input_ids): return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4) @@ -124,7 +195,6 @@ def test_decode_advances_trailing_text_offset_across_multiple_steps(): trailing_text = torch.arange(8, dtype=torch.float32).reshape(2, 4) state_tail = trailing_text last_hidden = torch.full((4,), 2.0, dtype=torch.float32) - tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32) meta = {"talker_text_offset": 0} seen_steps = [] @@ -135,7 +205,6 @@ def test_decode_advances_trailing_text_offset_across_multiple_steps(): text=["hello"], task_type=["CustomVoice"], hidden_states={"trailing_text": state_tail, "last": last_hidden}, - embed={"tts_pad": tts_pad}, meta=meta, _omni_is_prefill=False, _omni_num_computed_tokens=2, @@ -154,7 +223,8 @@ def test_decode_advances_trailing_text_offset_across_multiple_steps(): def test_decode_compacts_long_trailing_text_after_large_offset(): - model = _make_minimal_talker() + tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16) + model = _make_minimal_talker(tts_pad_embed=tts_pad) def fake_embed_input_ids(input_ids): return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4) @@ -162,7 +232,6 @@ def test_decode_compacts_long_trailing_text_after_large_offset(): model.embed_input_ids = fake_embed_input_ids trailing_text = torch.arange(130 * 4, dtype=torch.float32).reshape(130, 4) last_hidden = torch.full((4,), 2.0, dtype=torch.float32) - tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32) _, _, update = model.preprocess( input_ids=torch.tensor([123], dtype=torch.long), @@ -170,7 +239,6 @@ def test_decode_compacts_long_trailing_text_after_large_offset(): text=["hello"], task_type=["CustomVoice"], hidden_states={"trailing_text": trailing_text, "last": last_hidden}, - embed={"tts_pad": tts_pad}, meta={"talker_text_offset": 64}, _omni_is_prefill=False, _omni_num_computed_tokens=2, @@ -183,7 +251,8 @@ def test_decode_compacts_long_trailing_text_after_large_offset(): def test_decode_batch_preprocess_matches_decode_state_updates(): - model = _make_minimal_talker() + tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16) + model = _make_minimal_talker(tts_pad_embed=tts_pad) def fake_embed_input_ids(input_ids): return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4) @@ -193,7 +262,6 @@ def test_decode_batch_preprocess_matches_decode_state_updates(): trailing_b = torch.arange(8, dtype=torch.float32).reshape(2, 4) + 100 last_a = torch.full((4,), 2.0, dtype=torch.float32) last_b = torch.full((4,), 3.0, dtype=torch.float32) - tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32) out_ids, out_embeds, past_hidden, text_step, updates = model.preprocess_decode_batch( input_ids=torch.tensor([101, 202], dtype=torch.long), @@ -202,14 +270,12 @@ def test_decode_batch_preprocess_matches_decode_state_updates(): "text": ["hello"], "task_type": ["Base"], "hidden_states": {"trailing_text": trailing_a, "last": last_a}, - "embed": {"tts_pad": tts_pad}, "meta": {"talker_text_offset": 1}, }, { "text": ["world"], "task_type": ["CustomVoice"], "hidden_states": {"trailing_text": trailing_b, "last": last_b}, - "embed": {"tts_pad": tts_pad}, "meta": {"talker_text_offset": 2}, }, ], @@ -228,53 +294,49 @@ def test_decode_batch_preprocess_matches_decode_state_updates(): assert updates[1]["hidden_states"]["trailing_text"].numel() == 0 +def _stub_text_embedding(device_param: torch.nn.Parameter): + """Build a lambda that emulates ``nn.Embedding`` for the ``_device()`` helper. + + ``Qwen3TTSPromptEmbedsBuilder._device()`` calls + ``next(self._text_embedding.parameters()).device``, so the stub needs a + ``parameters()`` attribute *and* must be callable as an embedding. + """ + fn = lambda ids: torch.ones((*ids.shape, 4), device=ids.device) # noqa: E731 + fn.parameters = lambda: iter([device_param]) # type: ignore[attr-defined] + return fn + + def test_base_voice_clone_normalizes_ref_audio_once_for_ref_code_and_speaker(): - model = _make_minimal_talker() + builder = _make_minimal_builder() device_param = torch.nn.Parameter(torch.empty(0)) - model.parameters = lambda: iter([device_param]) - model.config = SimpleNamespace( - tts_bos_token_id=100, - tts_eos_token_id=101, - tts_pad_token_id=102, - ) - model.talker_config = SimpleNamespace( - codec_nothink_id=10, - codec_think_bos_id=11, - codec_think_eos_id=12, - codec_think_id=13, - codec_language_id={}, - codec_pad_id=7, - codec_bos_id=8, - num_code_groups=2, - spk_is_dialect={}, - ) + builder._text_embedding = _stub_text_embedding(device_param) + builder._text_projection = lambda embeds: embeds + builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device) class FakeTokenizer: def __call__(self, *_args, **_kwargs): return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)} - model._get_tokenizer = lambda: FakeTokenizer() - model.text_embedding = lambda ids: torch.ones((*ids.shape, 4), device=ids.device) - model.text_projection = lambda embeds: embeds - model.embed_input_ids = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device) - model._generate_icl_prompt = lambda **kwargs: ( + builder._text_tokenizer = FakeTokenizer() + + builder._generate_icl_prompt = lambda **kwargs: ( torch.ones((1, 2, 4), device=kwargs["ref_code"].device), torch.ones((1, 4), device=kwargs["ref_code"].device), ) normalize_calls = [] ref_audio = np.arange(1024, dtype=np.float32) - model._normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000) + builder.normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000) ref_audio_ids = [] - model._encode_ref_audio_batch = lambda wavs, sr, *, device: ( + builder._encode_ref_audio_batch_fn = lambda wavs, sr, *, device: ( [ref_audio_ids.append(id(w)) or torch.ones((2, 2), dtype=torch.long) for w in wavs] ) - model._extract_speaker_embedding = lambda wav, _sr: ( + builder.extract_speaker_embedding = lambda wav, _sr: ( ref_audio_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16) ) - _prompt, _trailing, _pad, ref_code_len, ref_code = model._build_prompt_embeds( + _prompt, _trailing, ref_code_len, ref_code = builder.build_prompt_embeds( task_type="Base", info_dict={ "text": ["hello"], @@ -291,11 +353,11 @@ def test_base_voice_clone_normalizes_ref_audio_once_for_ref_code_and_speaker(): def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate(): - model = _make_minimal_talker() + builder = _make_minimal_builder() wav1 = np.arange(2048, dtype=np.float32) wav2 = np.arange(3072, dtype=np.float32) normalize_calls = [] - model._normalize_ref_audio = lambda raw: ( + builder.normalize_ref_audio = lambda raw: ( normalize_calls.append(raw) or ( wav1 if raw == "a.wav" else wav2, @@ -312,7 +374,7 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate(): torch.full((3, 2), 22, dtype=torch.long), ] - model._encode_ref_audio_batch = _fake_encode_batch + builder._encode_ref_audio_batch_fn = _fake_encode_batch class FakeTextTokenizer: def __init__(self): @@ -323,7 +385,7 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate(): return {"input_ids": [[idx + 1, idx + 2, idx + 3] for idx, _ in enumerate(texts)]} text_tok = FakeTextTokenizer() - model._get_tokenizer = lambda: text_tok + builder._text_tokenizer = text_tok buf = { "r1": { "task_type": ["Base"], @@ -341,7 +403,7 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate(): }, } - model.preprocess_batch( + builder.preprocess_batch( req_ids=["r1", "r2"], model_intermediate_buffer=buf, device=torch.device("cpu"), @@ -353,30 +415,30 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate(): assert wavs[0] is wav1 assert wavs[1] is wav2 assert sr == 16000 - assert torch.equal(buf["r1"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((2, 2), 11)) - assert torch.equal(buf["r2"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((3, 2), 22)) - assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav1 - assert buf["r2"][_NORMALIZED_REF_AUDIO_KEY][0] is wav2 + assert torch.equal(buf["r1"]["codes"][PRECOMPUTED_REF_CODE_KEY], torch.full((2, 2), 11)) + assert torch.equal(buf["r2"]["codes"][PRECOMPUTED_REF_CODE_KEY], torch.full((3, 2), 22)) + assert buf["r1"][NORMALIZED_REF_AUDIO_KEY][0] is wav1 + assert buf["r2"][NORMALIZED_REF_AUDIO_KEY][0] is wav2 assert len(text_tok.calls) == 2 - assert torch.equal(buf["r1"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3])) - assert torch.equal(buf["r2"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4])) - assert torch.equal(buf["r1"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3])) - assert torch.equal(buf["r2"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4])) + assert torch.equal(buf["r1"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3])) + assert torch.equal(buf["r2"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4])) + assert torch.equal(buf["r1"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3])) + assert torch.equal(buf["r2"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4])) def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_without_speech_tokenizer(): - model = _make_minimal_talker() + builder = _make_minimal_builder() wav = np.arange(2048, dtype=np.float32) - model._normalize_ref_audio = lambda raw: (wav, 16000) - model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw( - AssertionError("singleton should not load speech tokenizer") + builder.normalize_ref_audio = lambda raw: (wav, 16000) + builder._encode_ref_audio_batch_fn = lambda wavs, sr, *, device: (_ for _ in ()).throw( + AssertionError("singleton should not call encode_ref_audio_batch") ) class FakeTextTokenizer: def __call__(self, texts, *, padding=False): return {"input_ids": [[7, 8, 9] for _ in texts]} - model._get_tokenizer = lambda: FakeTextTokenizer() + builder._text_tokenizer = FakeTextTokenizer() buf = { "r1": { "task_type": ["Base"], @@ -387,22 +449,27 @@ def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_wit } } - model.preprocess_batch( + builder.preprocess_batch( req_ids=["r1"], model_intermediate_buffer=buf, device=torch.device("cpu"), ) - assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav - assert torch.equal(buf["r1"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([7, 8, 9])) - assert torch.equal(buf["r1"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([7, 8, 9])) + assert buf["r1"][NORMALIZED_REF_AUDIO_KEY][0] is wav + assert torch.equal(buf["r1"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([7, 8, 9])) + assert torch.equal(buf["r1"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([7, 8, 9])) def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exists(): - model = _make_minimal_talker() - model._normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected")) - model._get_tokenizer = lambda: (_ for _ in ()).throw(AssertionError("tokenizer not expected")) - model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw( + builder = _make_minimal_builder() + builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected")) + + class _RaisingTokenizer: + def __call__(self, *_args, **_kwargs): + raise AssertionError("tokenizer not expected") + + builder._text_tokenizer = _RaisingTokenizer() + builder._encode_ref_audio_batch_fn = lambda *a, **kw: (_ for _ in ()).throw( AssertionError("speech tokenizer not expected") ) buf = { @@ -416,70 +483,53 @@ def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exi } } - model.preprocess_batch( + builder.preprocess_batch( req_ids=["r1"], model_intermediate_buffer=buf, device=torch.device("cpu"), ) - assert _PRECOMPUTED_TEXT_IDS_KEY not in buf["r1"] - assert _NORMALIZED_REF_AUDIO_KEY not in buf["r1"] + assert PRECOMPUTED_TEXT_IDS_KEY not in buf["r1"] + assert NORMALIZED_REF_AUDIO_KEY not in buf["r1"] def test_base_voice_clone_uses_batched_ref_code_without_serial_encode(): - model = _make_minimal_talker() + builder = _make_minimal_builder() device_param = torch.nn.Parameter(torch.empty(0)) - model.parameters = lambda: iter([device_param]) - model.config = SimpleNamespace( - tts_bos_token_id=100, - tts_eos_token_id=101, - tts_pad_token_id=102, - ) - model.talker_config = SimpleNamespace( - codec_nothink_id=10, - codec_think_bos_id=11, - codec_think_eos_id=12, - codec_think_id=13, - codec_language_id={}, - codec_pad_id=7, - codec_bos_id=8, - num_code_groups=2, - spk_is_dialect={}, - ) + builder._text_embedding = _stub_text_embedding(device_param) + builder._text_projection = lambda embeds: embeds + builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device) class FakeTokenizer: def __call__(self, *_args, **_kwargs): return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)} - model._get_tokenizer = lambda: FakeTokenizer() - model.text_embedding = lambda ids: torch.ones((*ids.shape, 4), device=ids.device) - model.text_projection = lambda embeds: embeds - model.embed_input_ids = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device) - model._generate_icl_prompt = lambda **kwargs: ( + builder._text_tokenizer = FakeTokenizer() + builder._generate_icl_prompt = lambda **kwargs: ( torch.ones((1, 2, 4), device=kwargs["ref_code"].device), torch.ones((1, 4), device=kwargs["ref_code"].device), ) ref_audio = np.arange(2048, dtype=np.float32) ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2) - model._normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("serial normalize not expected")) - model._encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw( + builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("serial normalize not expected")) + builder.encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw( AssertionError("serial encode not expected") ) speaker_wav_ids = [] - model._extract_speaker_embedding = lambda wav, _sr: ( + builder.extract_speaker_embedding = lambda wav, _sr: ( speaker_wav_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16) ) - _prompt, _trailing, _pad, ref_code_len, out_ref_code = model._build_prompt_embeds( + _prompt, _trailing, ref_code_len, out_ref_code = builder.build_prompt_embeds( task_type="Base", info_dict={ "text": ["hello"], "ref_audio": ["ref.wav"], "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1), "non_streaming_mode": [False], - "codes": {_PRECOMPUTED_REF_CODE_KEY: ref_code}, - _NORMALIZED_REF_AUDIO_KEY: (ref_audio, 16000), + "codes": {PRECOMPUTED_REF_CODE_KEY: ref_code}, + NORMALIZED_REF_AUDIO_KEY: (ref_audio, 16000), }, ) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py index c63622eee..9493a00f2 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py @@ -9,8 +9,8 @@ dimension and reports `len(ref_code) == num_codebooks` instead of import pytest -from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( - Qwen3TTSTalkerForConditionalGeneration, +from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + Qwen3TTSPromptEmbedsBuilder, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -38,7 +38,7 @@ def test_estimate_prompt_len_uses_full_ref_code_length() -> None: "language": ["English"], } - est = Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + est = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( additional_information=info, task_type="Base", tokenize_prompt=_fake_tokenize, @@ -68,7 +68,7 @@ def test_estimate_prompt_len_handles_1d_ref_code() -> None: "language": ["English"], } - est = Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + est = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( additional_information=info, task_type="Base", tokenize_prompt=_fake_tokenize, @@ -104,14 +104,14 @@ def test_estimate_prompt_len_uses_list_of_int_ref_ids_from_voice_clone_prompt() "language": ["English"], } - est_30 = Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + est_30 = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( additional_information=make_info(30), task_type="Base", tokenize_prompt=_fake_tokenize, codec_language_id={"english": 0}, spk_is_dialect=None, ) - est_50 = Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + est_50 = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( additional_information=make_info(50), task_type="Base", tokenize_prompt=_fake_tokenize, diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 458e40653..e2a4b5af8 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -660,8 +660,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int: """Estimate prompt length so the placeholder matches model-side embeddings.""" try: - from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( - Qwen3TTSTalkerForConditionalGeneration, + from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + Qwen3TTSPromptEmbedsBuilder, ) if self._tts_tokenizer is None: @@ -676,7 +676,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): hf_config = self.engine_client.model_config.hf_config talker_config = hf_config.talker_config task_type = (tts_params.get("task_type") or ["CustomVoice"])[0] - return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( additional_information=tts_params, task_type=task_type, tokenize_prompt=lambda t: self._tts_tokenizer(t, padding=False)["input_ids"], diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py new file mode 100644 index 000000000..ca756b794 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py @@ -0,0 +1,1543 @@ +"""Stand-alone builder for Qwen3-TTS AR talker prompt embeddings. + +Factors the prompt-construction logic out of +:class:`vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneration` +so alternative AR talker implementations that share the same backbone +embedding tables and codec / speaker encoders can reuse it without +duplicating the (substantial) prompt layout logic. + +The builder is intentionally *not* coupled to the talker class. The talker +constructs a builder, supplies its frozen embedding tables, projection +layers, codec embed lookups, speaker encoder, the precomputed pad-token +embedding buffer, and a callable that batch-encodes reference audio into +codec frames (so the builder doesn't have to know whether the talker +uses :class:`Qwen3TTSTokenizer` or :class:`Qwen3TTSTokenizerV2Encoder` +under the hood). The single public entry point +:meth:`Qwen3TTSPromptEmbedsBuilder.build_prompt_embeds` mirrors the +official ``Qwen3TTSForConditionalGeneration`` prompt layout for the three +supported task types (``CustomVoice``, ``VoiceDesign``, ``Base``) and +both streaming / non-streaming modes. +""" + +from __future__ import annotations + +import base64 +import hashlib +import io +from collections import OrderedDict +from collections.abc import Callable, Mapping, Sequence +from functools import lru_cache +from typing import TYPE_CHECKING, Any +from urllib.parse import urlparse + +import numpy as np +import soundfile as sf +import torch +import torch.nn as nn +from transformers import AutoTokenizer +from vllm.logger import init_logger +from vllm.multimodal.audio import AudioResampler + +from vllm_omni.utils.audio import mel_filter_bank + +if TYPE_CHECKING: + from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSTalkerConfig + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Constants shared with the talker's preprocess pipeline. +# +# These keys mark precomputed artifacts that the batched preprocess pass +# stashes into ``info_dict`` so the per-request prompt build can skip +# redundant tokenization / codec encoding work. +# --------------------------------------------------------------------------- + +PRECOMPUTED_REF_CODE_KEY = "precomputed_ref" +NORMALIZED_REF_AUDIO_KEY = "_qwen3_tts_normalized_ref_audio" +REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key" +PRECOMPUTED_TEXT_IDS_KEY = "_qwen3_tts_text_ids" +PRECOMPUTED_REF_IDS_KEY = "_qwen3_tts_ref_ids" + + +# --------------------------------------------------------------------------- +# Chat-template helpers (must match the official HuggingFace reference). +# --------------------------------------------------------------------------- + + +def build_assistant_text(text: str) -> str: + return f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + + +def build_ref_text(text: str) -> str: + return f"<|im_start|>assistant\n{text}<|im_end|>\n" + + +def build_instruct_text(instruct: str) -> str: + return f"<|im_start|>user\n{instruct}<|im_end|>\n" + + +# --------------------------------------------------------------------------- +# Type-coercion helpers. All accept the loose shapes that survive +# msgspec IPC ``additional_information`` round-trips (scalars, single- +# element lists, ndarrays, tensors). +# --------------------------------------------------------------------------- + + +def first_value(value: object, default: object = None) -> object: + if isinstance(value, list): + return value[0] if value else default + return value if value is not None else default + + +def coerce_ref_code_tensor(value: object, *, device: torch.device) -> torch.Tensor | None: + value = first_value(value) + if value is None: + return None + if isinstance(value, torch.Tensor): + ref_code = value + elif isinstance(value, np.ndarray): + ref_code = torch.from_numpy(value) + elif isinstance(value, list) and value: + ref_code = torch.as_tensor(value, dtype=torch.long) + else: + return None + if ref_code.ndim == 3: + ref_code = ref_code[0] + if ref_code.ndim != 2 or ref_code.numel() == 0: + return None + return ref_code.to(device=device, dtype=torch.long).contiguous() + + +def coerce_token_ids(value: object, *, device: torch.device) -> torch.Tensor | None: + value = first_value(value) + if value is None: + return None + if isinstance(value, torch.Tensor): + ids = value + elif isinstance(value, np.ndarray): + ids = torch.from_numpy(value) + elif isinstance(value, list) and value and all(isinstance(v, (int, np.integer)) for v in value): + ids = torch.tensor(value, dtype=torch.long) + else: + return None + if ids.ndim == 1: + ids = ids.unsqueeze(0) + if ids.ndim != 2 or ids.numel() == 0: + return None + return ids.to(device=device, dtype=torch.long).contiguous() + + +# --------------------------------------------------------------------------- +# Mel-spectrogram utilities used by the speaker encoder front-end. +# Module-level LRU caches for the (sr, n_fft, n_mels, ...) -> mel_basis and +# win_size -> hann_window tensors avoid rebuilding them on every request. +# --------------------------------------------------------------------------- + + +def _dynamic_range_compression(x: torch.Tensor, c: float = 1, clip_val: float = 1e-5) -> torch.Tensor: + return torch.log(torch.clamp(x, min=clip_val) * c) + + +@lru_cache(maxsize=8) +def _cached_mel_filter_bank(sampling_rate: int, n_fft: int, n_mels: int, fmin: int, fmax: int | None) -> torch.Tensor: + return mel_filter_bank(sr=sampling_rate, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax) + + +@lru_cache(maxsize=8) +def _cached_hann_window(win_size: int) -> torch.Tensor: + return torch.hann_window(win_size) + + +def mel_spectrogram( + y: torch.Tensor, + n_fft: int, + num_mels: int, + sampling_rate: int, + hop_size: int, + win_size: int, + fmin: int, + fmax: int | None = None, + center: bool = False, + mel_basis: torch.Tensor | None = None, + hann_window: torch.Tensor | None = None, +) -> torch.Tensor: + """Mel spectrogram via torch STFT and a (cached) mel filterbank.""" + if torch.min(y) < -1.0: + logger.warning("Min value of input waveform signal is %s", torch.min(y)) + if torch.max(y) > 1.0: + logger.warning("Max value of input waveform signal is %s", torch.max(y)) + device = y.device + if mel_basis is None: + mel_basis = _cached_mel_filter_bank(sampling_rate, n_fft, num_mels, fmin, fmax).to(device) + elif mel_basis.device != device: + mel_basis = mel_basis.to(device) + if hann_window is None: + hann_window = _cached_hann_window(win_size).to(device) + elif hann_window.device != device: + hann_window = hann_window.to(device) + padding = (n_fft - hop_size) // 2 + y = torch.nn.functional.pad(y.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1) + spec = torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=hann_window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9) + mel_spec = torch.matmul(mel_basis, spec) + return _dynamic_range_compression(mel_spec) + + +# --------------------------------------------------------------------------- +# Local audio loading helpers (URL / base64 / local path). +# --------------------------------------------------------------------------- + + +def _is_probably_base64(s: str) -> bool: + if s.startswith("data:audio"): + return True + if ("/" not in s and "\\" not in s) and len(s) > 256: + return True + return False + + +def _is_url(s: str) -> bool: + try: + u = urlparse(s) + if u.scheme in ("http", "https"): + return bool(u.netloc) + return u.scheme == "file" + except Exception: + return False + + +def _decode_base64_to_wav_bytes(b64: str) -> bytes: + if "," in b64 and b64.strip().startswith("data:"): + b64 = b64.split(",", 1)[1] + return base64.b64decode(b64) + + +def load_audio_to_np(x: str) -> tuple[np.ndarray, int]: + """Load mono float32 audio from a local path, URL, or base64 data URI. + + Uses upstream vLLM's MediaConnector for http(s) URLs and ``file:`` URIs + with unrestricted local access (offline inference is trusted). + """ + from vllm.multimodal.media.audio import load_audio + + if _is_url(x): + from vllm.multimodal.media import MediaConnector + + connector = MediaConnector(allowed_local_media_path="/") + audio, sr = connector.fetch_audio(x) + elif _is_probably_base64(x): + wav_bytes = _decode_base64_to_wav_bytes(x) + with io.BytesIO(wav_bytes) as f: + audio, sr = sf.read(f, dtype="float32", always_2d=False) + else: + audio, sr = load_audio(x, sr=None, mono=True) + + if isinstance(audio, np.ndarray) and audio.ndim > 1: + audio = np.mean(audio, axis=-1) + + return np.asarray(audio, dtype=np.float32), int(sr) + + +# --------------------------------------------------------------------------- +# Prompt-embeds builder. +# --------------------------------------------------------------------------- + + +class Qwen3TTSPromptEmbedsBuilder: + """Build prefill prompt embeddings for the Qwen3-TTS AR talker. + + Stand-alone wrapper around the embedding layers, encoders, and speaker + cache required to assemble a talker prompt. Mirrors the layout of the + official HuggingFace ``Qwen3TTSForConditionalGeneration`` pipeline for + ``task_type in {CustomVoice, VoiceDesign, Base}`` and both streaming / + non-streaming modes. + + The builder does **not** own the speech / codec tokenizer used to turn + raw reference audio into codec frames. The caller passes in a + ``encode_ref_audio_batch`` callable so the same builder works against + either the legacy :class:`Qwen3TTSTokenizer` HF wrapper or the native + :class:`Qwen3TTSTokenizerV2Encoder` exposed on the talker. + + Args: + config: Top-level Qwen3-TTS config; supplies ``tts_*_token_id`` and + ``speaker_encoder_config``. + talker_config: Talker sub-config; supplies the codec control-token + ids, ``num_code_groups``, ``spk_id``, ``spk_is_dialect`` and + ``codec_language_id`` maps. + model_path: HuggingFace repo or local path used to lazily load the + text tokenizer. + text_embedding: Text-side embedding table. + text_projection: Callable projecting text-side embeddings into the + talker hidden size. + codec_embed: Callable performing codec-side embedding lookup + (normally the talker backbone's ``embed_input_ids``). + residual_code_embeddings: Callable returning the per-codebook + embedding modules used by the code predictor; called lazily so + the predictor doesn't need to be constructed before the builder. + speaker_encoder: ECAPA-TDNN speaker encoder used by the Base task. + tts_pad_embed: Pre-computed ``tts_pad`` embedding (shape + ``[1, hidden]``) reused at every prefill to avoid recomputing. + encode_ref_audio_batch: Callable ``(wavs, sr, *, device) -> list[Tensor]`` + that returns one ``[T, num_quantizers]`` long tensor per input + waveform. Single-waveform encoding is implemented in terms of + this callable. + speaker_cache: Optional :class:`SpeakerEmbeddingCache`. When + provided, Base ``CustomVoice``-style results are cached by + ``(speaker_name, mode, created_at)``. + ref_audio_artifact_cache_max_entries: Upper bound on the SHA1 + keyed LRU cache of per-ref-audio ``(ref_code, ref_spk_embedding)`` + artifacts. ``0`` disables the cache. + """ + + def __init__( + self, + *, + config: Qwen3TTSConfig, + talker_config: Qwen3TTSTalkerConfig, + model_path: str, + text_embedding: nn.Module, + text_projection: Callable[[torch.Tensor], torch.Tensor], + codec_embed: Callable[[torch.Tensor], torch.Tensor], + residual_code_embeddings: Callable[[], Sequence[Callable[[torch.Tensor], torch.Tensor]]], + speaker_encoder: nn.Module, + tts_pad_embed: torch.Tensor, + encode_ref_audio_batch: Callable[..., list[torch.Tensor]], + speaker_cache: Any | None = None, + ref_audio_artifact_cache_max_entries: int = 256, + ): + self._config = config + self._talker_config = talker_config + self._model_path = model_path + self._text_embedding = text_embedding + self._text_projection = text_projection + self._codec_embed = codec_embed + self._residual_code_embeddings = residual_code_embeddings + self._speaker_encoder = speaker_encoder + self._tts_pad_embed_buffer = tts_pad_embed + self._encode_ref_audio_batch_fn = encode_ref_audio_batch + self._speaker_cache = speaker_cache + + self._text_tokenizer: Any | None = None + + self._ref_audio_artifact_cache_max_entries = int(ref_audio_artifact_cache_max_entries) + self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict() + + # Bounded LRU; caller-supplied orig_sr can otherwise grow this without limit. + self._resampler_cache: OrderedDict[tuple[int, int], AudioResampler] = OrderedDict() + self._resampler_cache_max = 16 + + # -------------------- device / buffer helpers -------------------- + + def _device(self) -> torch.device: + # text_embedding is always present and lives on the talker device. + return next(self._text_embedding.parameters()).device + + def _pad_embed(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """Return the pad-token embedding on the requested device/dtype, shape ``[1, 1, H]``.""" + return self._tts_pad_embed_buffer.to(device=device, dtype=dtype).reshape(1, 1, -1) + + def _get_resampler(self, orig_sr: int, target_sr: int) -> AudioResampler: + key = (int(orig_sr), int(target_sr)) + cache = self._resampler_cache + if key in cache: + cache.move_to_end(key) + return cache[key] + cache[key] = AudioResampler(target_sr=int(target_sr)) + if len(cache) > self._resampler_cache_max: + cache.popitem(last=False) + return cache[key] + + # -------------------- text tokenizer (lazy) -------------------- + + def get_text_tokenizer(self): + if self._text_tokenizer is None: + import transformers + + kwargs: dict[str, Any] = dict(trust_remote_code=True, use_fast=True) + if transformers.__version__ < "5": + kwargs["fix_mistral_regex"] = True + tok = AutoTokenizer.from_pretrained(self._model_path, **kwargs) + tok.padding_side = "left" + self._text_tokenizer = tok + return self._text_tokenizer + + # -------------------- ref-audio handling -------------------- + + def load_audio_to_np(self, x: str) -> tuple[np.ndarray, int]: + return load_audio_to_np(x) + + def normalize_ref_audio(self, ref_audio: object) -> tuple[np.ndarray, int]: + """Coerce a serialized ref_audio payload into ``(wav_np, sample_rate)``. + + ``additional_information`` may serialize ``(wav, sr)`` into (nested) + lists across processes, so we accept and unwrap a variety of shapes. + """ + if isinstance(ref_audio, str): + return self.load_audio_to_np(ref_audio) + + def _is_sr(x: object) -> bool: + try: + v = int(x) # type: ignore[arg-type] + except Exception: + return False + return 1_000 <= v <= 200_000 + + def _is_number_sequence(xs: list[object]) -> bool: + if not xs: + return False + for v in xs[:8]: + if not isinstance(v, (int, float, np.number)): + return False + return True + + wav_candidates: list[object] = [] + sr_candidates: list[int] = [] + + def _summarize(obj: object, depth: int = 0) -> str: + if depth > 2: + if isinstance(obj, (int, np.integer)): + return f"int({int(obj)})" + return type(obj).__name__ + if obj is None: + return "None" + if isinstance(obj, str): + if len(obj) <= 16: + return f"str({obj!r})" + return f"str(len={len(obj)})" + if isinstance(obj, (int, float, np.number)): + return f"{type(obj).__name__}({obj})" + if isinstance(obj, np.ndarray): + return f"ndarray(shape={obj.shape}, dtype={obj.dtype})" + if isinstance(obj, torch.Tensor): + return f"Tensor(shape={tuple(obj.shape)}, dtype={obj.dtype}, device={obj.device})" + if isinstance(obj, dict): + keys = list(obj.keys()) + return f"dict(keys={keys[:8]})" + if isinstance(obj, (tuple, list)): + items = list(obj) + head = ", ".join(_summarize(x, depth + 1) for x in items[:3]) + return f"{type(obj).__name__}(len={len(items)}; head=[{head}])" + return f"{type(obj).__name__}" + + def _scan(obj: object, depth: int = 0) -> None: + if depth > 4: + return + if obj is None: + return + if _is_sr(obj): + sr_candidates.append(int(obj)) # type: ignore[arg-type] + return + if isinstance(obj, np.ndarray) and obj.size > 0: + wav_candidates.append(obj) + return + if isinstance(obj, torch.Tensor) and obj.numel() > 0: + wav_candidates.append(obj) + return + if isinstance(obj, dict): + # Inlined ndarray/tensor payloads from the input processor. + if obj.get("__ndarray__") and "data" in obj and "dtype" in obj and "shape" in obj: + try: + data = obj["data"] + dtype = obj["dtype"] + shape = obj["shape"] + if isinstance(data, (bytes, bytearray, memoryview)): + arr = np.frombuffer(data, dtype=dtype).reshape(shape) + if arr.size > 0: + wav_candidates.append(arr) + return + except Exception: + pass + if obj.get("__tensor__") and "data" in obj and "dtype" in obj and "shape" in obj: + try: + data = obj["data"] + dtype = obj["dtype"] + shape = obj["shape"] + if isinstance(data, (bytes, bytearray, memoryview)): + # Stored as raw CPU bytes; interpret as numpy for audio. + np_dtype = np.dtype(dtype) + arr = np.frombuffer(data, dtype=np_dtype).reshape(shape) + if arr.size > 0: + wav_candidates.append(arr) + return + except Exception: + pass + wav_obj = obj.get("array") or obj.get("wav") or obj.get("audio") + sr_obj = obj.get("sampling_rate") or obj.get("sr") or obj.get("sample_rate") + if wav_obj is not None: + _scan(wav_obj, depth + 1) + if sr_obj is not None: + _scan(sr_obj, depth + 1) + return + if isinstance(obj, (tuple, list)): + obj_list = list(obj) + # Unwrap singleton nesting ([[wav, sr]]). + while isinstance(obj_list, list) and len(obj_list) == 1: + inner = obj_list[0] + if isinstance(inner, np.ndarray) and inner.size > 0: + wav_candidates.append(inner) + return + if isinstance(inner, torch.Tensor) and inner.numel() > 0: + wav_candidates.append(inner) + return + if isinstance(inner, dict): + _scan(inner, depth + 1) + return + if isinstance(inner, (tuple, list)): + obj_list = list(inner) # type: ignore[list-item] + continue + break + + # If the *unwrapped* list is a long list of numbers, treat it as waveform. + if len(obj_list) >= 512 and _is_number_sequence(obj_list): + wav_candidates.append(obj_list) + return + + # Otherwise, recurse into elements (but avoid descending into huge numeric lists). + for item in obj_list: + if isinstance(item, list) and len(item) >= 512 and _is_number_sequence(item): # type: ignore[arg-type] + wav_candidates.append(item) + continue + _scan(item, depth + 1) + return + + _scan(ref_audio) + if not sr_candidates: + raise TypeError(f"ref_audio missing sample_rate: {_summarize(ref_audio)}") + sr = int(sr_candidates[0]) + + def _wav_len(x: object) -> int: + try: + if isinstance(x, np.ndarray): + return int(x.size) + if isinstance(x, torch.Tensor): + return int(x.numel()) + if isinstance(x, list): + return int(len(x)) + except Exception: + pass + return 0 + + if not wav_candidates: + raise TypeError(f"ref_audio missing waveform: {_summarize(ref_audio)}") + wav_obj = max(wav_candidates, key=_wav_len) + + def _to_np(x: object) -> np.ndarray: + if isinstance(x, np.ndarray): + return x.astype(np.float32).reshape(-1) + if isinstance(x, torch.Tensor): + return x.detach().to("cpu").float().contiguous().numpy().reshape(-1) + if isinstance(x, dict) and x.get("__ndarray__") and "data" in x and "dtype" in x and "shape" in x: + data = x["data"] + dtype = x["dtype"] + shape = x["shape"] + if isinstance(data, (bytes, bytearray, memoryview)): + return np.frombuffer(data, dtype=dtype).reshape(shape).astype(np.float32).reshape(-1) + if isinstance(x, list): + # list of numbers + if len(x) >= 2 and _is_number_sequence(x): # type: ignore[arg-type] + return np.asarray(x, dtype=np.float32).reshape(-1) + # list of chunks + parts: list[np.ndarray] = [] + for part in x: + if isinstance(part, (np.ndarray, torch.Tensor, list)): + parts.append(_to_np(part)) + if parts: + return np.concatenate(parts, axis=0) + raise TypeError(f"Unsupported waveform type: {type(x)}") + + wav_np = _to_np(wav_obj) + if wav_np.size < 1024: + raise ValueError(f"ref_audio waveform too short: {wav_np.size} samples") + return wav_np, sr + + # -------------------- ref-audio artifact cache -------------------- + + @staticmethod + def make_ref_audio_cache_key(wav: np.ndarray, sr: int) -> str: + wav_f32 = wav.astype(np.float32, copy=False).reshape(-1) + h = hashlib.sha1() + h.update(int(sr).to_bytes(4, byteorder="little", signed=False)) + h.update(int(wav_f32.size).to_bytes(8, byteorder="little", signed=False)) + h.update(wav_f32.tobytes(order="C")) + return h.hexdigest() + + def get_ref_audio_artifacts(self, cache_key: str) -> dict[str, torch.Tensor | bool] | None: + if self._ref_audio_artifact_cache_max_entries <= 0: + return None + entry = self._ref_audio_artifact_cache.get(cache_key) + if entry is None: + return None + self._ref_audio_artifact_cache.move_to_end(cache_key, last=True) + return entry + + def put_ref_audio_artifacts( + self, + cache_key: str, + *, + ref_code: torch.Tensor | None = None, + ref_spk_embedding: torch.Tensor | None = None, + ) -> None: + if self._ref_audio_artifact_cache_max_entries <= 0 or not cache_key: + return + entry = self._ref_audio_artifact_cache.get(cache_key) + if entry is None: + entry = {} + if isinstance(ref_code, torch.Tensor): + entry["ref_code"] = ref_code.detach().to("cpu", dtype=torch.long).contiguous() + if isinstance(ref_spk_embedding, torch.Tensor): + entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=torch.bfloat16).reshape(-1) + if not entry: + return + self._ref_audio_artifact_cache[cache_key] = entry + self._ref_audio_artifact_cache.move_to_end(cache_key, last=True) + while len(self._ref_audio_artifact_cache) > self._ref_audio_artifact_cache_max_entries: + self._ref_audio_artifact_cache.popitem(last=False) + + # -------------------- speaker encoder / codec encoder -------------------- + + def extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor: + # vLLM workers do not automatically move arbitrary torch.nn.Modules to + # CUDA. Ensure the speaker encoder is on the same device/dtype as the + # main model before running it. + dev = self._device() + try: + spk_param = next(self._speaker_encoder.parameters()) + if spk_param.device != dev or spk_param.dtype != torch.bfloat16: + self._speaker_encoder.to(device=dev, dtype=torch.bfloat16) + except StopIteration: + pass + + # Resample to 24kHz for speaker encoder. + target_sr = int(getattr(self._config.speaker_encoder_config, "sample_rate", 24000)) + if sr != target_sr: + resampler = self._get_resampler(int(sr), target_sr) + wav = resampler.resample(wav.astype(np.float32), orig_sr=int(sr)) + sr = target_sr + + # Follow official implementation: mel_spectrogram expects 24kHz. Move + # the waveform first so STFT/mel computation stays on the model device + # instead of materializing a CPU mel tensor and copying it per request. + wav_tensor = torch.from_numpy(wav).to(device=dev, dtype=torch.float32).unsqueeze(0) + mels = mel_spectrogram( + wav_tensor, + n_fft=1024, + num_mels=128, + sampling_rate=24000, + hop_size=256, + win_size=1024, + fmin=0, + fmax=12000, + ).transpose(1, 2) + spk = self._speaker_encoder(mels.to(dtype=torch.bfloat16))[0] + return spk.to(dtype=torch.bfloat16) + + def encode_ref_audio_batch( + self, + wavs: list[np.ndarray], + sr: int, + *, + device: torch.device, + ) -> list[torch.Tensor]: + """Delegate to the talker-provided batched ref-audio encoder.""" + return self._encode_ref_audio_batch_fn(wavs, int(sr), device=device) + + def encode_ref_audio_to_code(self, wav: np.ndarray, sr: int) -> torch.Tensor: + """Single-waveform convenience wrapper around :meth:`encode_ref_audio_batch`.""" + codes = self.encode_ref_audio_batch([wav], int(sr), device=self._device()) + if not codes: + raise ValueError("encode_ref_audio_batch returned no codes") + return codes[0] + + # -------------------- batched preprocess (cross-request) -------------------- + + @staticmethod + def _voice_clone_prompt_dict(raw: object) -> dict[str, object] | None: + raw = first_value(raw) + if isinstance(raw, dict): + return raw + if isinstance(raw, list) and raw and isinstance(raw[0], dict): + return raw[0] + return None + + @staticmethod + def _has_ref_code_like(value: object) -> bool: + value = first_value(value) + if isinstance(value, torch.Tensor): + return value.numel() > 0 and value.ndim in (2, 3) + if isinstance(value, np.ndarray): + return value.size > 0 and value.ndim in (2, 3) + if isinstance(value, list): + return bool(value) + return False + + @staticmethod + def _needs_initial_prompt_preprocess(info_dict: dict[str, Any]) -> bool: + embed = info_dict.get("embed") + return not (isinstance(embed, dict) and "prefill" in embed) + + def _needs_batched_ref_code(self, info_dict: dict[str, Any]) -> bool: + if not self._needs_initial_prompt_preprocess(info_dict): + return False + if first_value(info_dict.get("task_type"), "CustomVoice") != "Base": + return False + codes = info_dict.get("codes") + if isinstance(codes, dict) and self._has_ref_code_like(codes.get(PRECOMPUTED_REF_CODE_KEY)): + return False + + voice_clone_prompt = self._voice_clone_prompt_dict(info_dict.get("voice_clone_prompt")) + if voice_clone_prompt is not None: + if self._has_ref_code_like(voice_clone_prompt.get("ref_code")): + return False + icl_flag = first_value(voice_clone_prompt.get("icl_mode")) + if isinstance(icl_flag, bool) and not icl_flag: + return False + + xvec_only = first_value(info_dict.get("x_vector_only_mode"), False) + if isinstance(xvec_only, bool) and xvec_only: + return False + + ref_audio_list = info_dict.get("ref_audio") + return isinstance(ref_audio_list, list) and bool(ref_audio_list) + + @torch.inference_mode() + def preprocess_batch( + self, + *, + req_ids: list[str], + model_intermediate_buffer: dict[str, dict[str, Any]], + device: torch.device, + ) -> None: + """Batch Base voice-clone ref-audio codec extraction for current prefill requests. + + Tokenizes assistant / ref text and runs the talker-supplied + ``encode_ref_audio_batch`` callable once per sample-rate group, + stashing the precomputed artifacts under + :data:`PRECOMPUTED_TEXT_IDS_KEY`, :data:`PRECOMPUTED_REF_IDS_KEY`, + :data:`PRECOMPUTED_REF_CODE_KEY`, and :data:`NORMALIZED_REF_AUDIO_KEY` + so the per-request :meth:`build_prompt_embeds` call can skip the + equivalent serial work. + """ + pending_text: list[tuple[dict[str, Any], str]] = [] + pending_ref_text: list[tuple[dict[str, Any], str]] = [] + groups: dict[int, list[tuple[dict[str, Any], np.ndarray, int]]] = {} + for req_id in req_ids: + info_dict = model_intermediate_buffer.get(req_id) + if not isinstance(info_dict, dict): + continue + + if ( + self._needs_initial_prompt_preprocess(info_dict) + and first_value(info_dict.get("task_type"), "CustomVoice") == "Base" + ): + if PRECOMPUTED_TEXT_IDS_KEY not in info_dict: + text = first_value(info_dict.get("text"), "") + if isinstance(text, str): + pending_text.append((info_dict, build_assistant_text(text))) + + if PRECOMPUTED_REF_IDS_KEY not in info_dict: + ref_text = first_value(info_dict.get("ref_text")) + if isinstance(ref_text, str) and ref_text.strip(): + pending_ref_text.append((info_dict, build_ref_text(ref_text))) + + if not self._needs_batched_ref_code(info_dict): + continue + ref_audio_list = info_dict.get("ref_audio") + if not isinstance(ref_audio_list, list) or not ref_audio_list: + continue + try: + wav, sr = self.normalize_ref_audio(ref_audio_list[0]) + except Exception: + # Keep the original per-request path responsible for surfacing + # invalid ref_audio errors with its existing messages. + continue + info_dict[NORMALIZED_REF_AUDIO_KEY] = (wav, sr) + cache_key = self.make_ref_audio_cache_key(wav, sr) + info_dict[REF_AUDIO_CACHE_KEY] = cache_key + cached = self.get_ref_audio_artifacts(cache_key) + if cached is not None: + cached_ref_code = coerce_ref_code_tensor(cached.get("ref_code"), device=device) + if isinstance(cached_ref_code, torch.Tensor): + info_dict.setdefault("codes", {})[PRECOMPUTED_REF_CODE_KEY] = cached_ref_code + continue + groups.setdefault(int(sr), []).append((info_dict, wav, int(sr))) + + if pending_text or pending_ref_text: + tok_text = self.get_text_tokenizer() + for items, key in ( + (pending_text, PRECOMPUTED_TEXT_IDS_KEY), + (pending_ref_text, PRECOMPUTED_REF_IDS_KEY), + ): + if not items: + continue + try: + tokenized = tok_text([text for _, text in items], padding=False) + except Exception as exc: + logger.debug("Qwen3-TTS batched text tokenization failed; falling back to serial path: %s", exc) + continue + input_ids = tokenized.get("input_ids") if isinstance(tokenized, dict) else None + if not isinstance(input_ids, list) or len(input_ids) != len(items): + continue + for (info_dict, _), ids in zip(items, input_ids, strict=True): + if isinstance(ids, list) and ids: + info_dict[key] = torch.tensor(ids, dtype=torch.long) + + groups = {sr: items for sr, items in groups.items() if len(items) >= 2} + if not groups: + return + + for sr, items in groups.items(): + wavs = [wav for _, wav, _ in items] + try: + ref_codes = self.encode_ref_audio_batch(wavs, int(sr), device=device) + except Exception as exc: + logger.debug("Qwen3-TTS batched ref_code encode failed; falling back to serial path: %s", exc) + continue + for (info_dict, wav, item_sr), ref_code in zip(items, ref_codes, strict=True): + ref_code_t = coerce_ref_code_tensor(ref_code, device=device) + if ref_code_t is None: + continue + info_dict.setdefault("codes", {})[PRECOMPUTED_REF_CODE_KEY] = ref_code_t + info_dict[NORMALIZED_REF_AUDIO_KEY] = (wav, item_sr) + cache_key = info_dict.get(REF_AUDIO_CACHE_KEY) + if isinstance(cache_key, str) and cache_key: + self.put_ref_audio_artifacts(cache_key, ref_code=ref_code_t) + + # -------------------- ICL prompt assembly -------------------- + + def _generate_icl_prompt( + self, + *, + text_id: torch.Tensor, + ref_id: torch.Tensor, + ref_code: torch.Tensor, + tts_pad_embed: torch.Tensor, + tts_eos_embed: torch.Tensor, + non_streaming_mode: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Port of the official ``Qwen3TTSForConditionalGeneration.generate_icl_prompt``.""" + text_embed = self._text_projection(self._text_embedding(torch.cat([ref_id, text_id], dim=-1))) + text_embed = torch.cat([text_embed, tts_eos_embed], dim=1) + + # codec embed (codec bos + codec) 1 T2 D + residual_embeds = self._residual_code_embeddings() + codec_embed: list[torch.Tensor] = [] + for i in range(int(self._talker_config.num_code_groups)): + if i == 0: + codec_embed.append(self._codec_embed(ref_code[:, :1])) + else: + codec_embed.append(residual_embeds[i - 1](ref_code[:, i : i + 1])) + codec_embed_sum = torch.cat(codec_embed, dim=1).sum(1).unsqueeze(0) # [1,T,H] + codec_embed_sum = torch.cat( + [ + self._codec_embed( + torch.tensor([[self._talker_config.codec_bos_id]], device=codec_embed_sum.device, dtype=torch.long) + ), + codec_embed_sum, + ], + dim=1, + ) + + text_lens = int(text_embed.shape[1]) + codec_lens = int(codec_embed_sum.shape[1]) + if non_streaming_mode: + # Official non-streaming mode: append the full text conditioning in + # prefill, and use PAD in decode steps. + icl_input_embed = text_embed + self._codec_embed( + torch.tensor( + [[self._talker_config.codec_pad_id] * text_lens], + device=codec_embed_sum.device, + dtype=torch.long, + ) + ) + icl_input_embed = torch.cat([icl_input_embed, codec_embed_sum + tts_pad_embed], dim=1) + return icl_input_embed, tts_pad_embed + if text_lens > codec_lens: + return text_embed[:, :codec_lens] + codec_embed_sum, text_embed[:, codec_lens:] + text_embed = torch.cat([text_embed] + [tts_pad_embed] * (codec_lens - text_lens), dim=1) + return text_embed + codec_embed_sum, tts_pad_embed + + # -------------------- main entry point -------------------- + + def build_prompt_embeds( + self, + *, + task_type: str, + info_dict: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor, int | None, torch.Tensor | None]: + """Build the prefill prompt embeddings for a single TTS request. + + Args: + task_type: One of ``"CustomVoice"``, ``"VoiceDesign"``, ``"Base"``. + info_dict: Per-request ``additional_information`` plus any + precomputed artifacts stashed by the batched preprocess pass + (``PRECOMPUTED_TEXT_IDS_KEY``, ``PRECOMPUTED_REF_IDS_KEY``, + ``PRECOMPUTED_REF_CODE_KEY``, ``NORMALIZED_REF_AUDIO_KEY``). + Consumed keys are popped so the caller doesn't re-feed them + on subsequent prefill chunks. + + Returns: + ``(talker_prompt, trailing_text_hidden, ref_code_len, ref_code)``: + + * ``talker_prompt``: ``[prompt_len, hidden]`` prefill embedding. + * ``trailing_text_hidden``: ``[T, hidden]`` queue of text-side + embeddings consumed one-per-decode-step (streaming mode) or a + single pad row (non-streaming). + * ``ref_code_len``: Number of ref codec frames consumed when the + Base task ran in in-context mode (``None`` otherwise). + * ``ref_code``: The ``[T, Q]`` ref codec tensor used for ICL + (``None`` outside Base / ICL). + """ + config = self._config + talker_config = self._talker_config + text_embedding = self._text_embedding + text_projection = self._text_projection + codec_embed = self._codec_embed + + text = (info_dict.get("text") or [""])[0] + language = (info_dict.get("language") or ["Auto"])[0] + non_streaming_mode_val = info_dict.get("non_streaming_mode") + if isinstance(non_streaming_mode_val, list): + non_streaming_mode_raw = non_streaming_mode_val[0] if non_streaming_mode_val else None + else: + non_streaming_mode_raw = non_streaming_mode_val + if isinstance(non_streaming_mode_raw, bool): + non_streaming_mode = non_streaming_mode_raw + else: + # Match official inference defaults: + # - CustomVoice/VoiceDesign: non_streaming_mode=True + # - Base: non_streaming_mode=False + non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign") + + # Text ids for assistant template (always). + tok = self.get_text_tokenizer() + dev = self._device() + input_ids = coerce_token_ids(info_dict.pop(PRECOMPUTED_TEXT_IDS_KEY, None), device=dev) + if input_ids is None: + input_ids = tok(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"].to(device=dev) + + # Optional instruct prefix. + instruct = (info_dict.get("instruct") or [""])[0] + instruct_embed = None + if isinstance(instruct, str) and instruct.strip(): + instruct_ids = tok(build_instruct_text(instruct), return_tensors="pt", padding=False)["input_ids"].to( + device=input_ids.device + ) + instruct_embed = text_projection(text_embedding(instruct_ids)) + + # tts special token embeds (projected into talker hidden). + # ``tts_pad_embed`` is precomputed (request-independent), so we only + # need bos/eos here. + tts_tokens = torch.tensor( + [[config.tts_bos_token_id, config.tts_eos_token_id]], + device=input_ids.device, + dtype=input_ids.dtype, + ) + tts_bos_embed, tts_eos_embed = text_projection(text_embedding(tts_tokens)).chunk(2, dim=1) + tts_pad_embed = self._pad_embed(input_ids.device, tts_bos_embed.dtype) + + # Codec prefill tags. + language_id = None + if isinstance(language, str) and language.lower() != "auto": + language_id = talker_config.codec_language_id.get(language.lower()) + # Match official dialect override: + # If language is Chinese/Auto and the selected speaker is a dialect voice, + # set language_id to that dialect to improve code generation stability. + if language_id is None and isinstance(language, str) and language.lower() in ("chinese", "auto"): + speaker_for_dialect = None + if task_type == "CustomVoice": + speaker_for_dialect = (info_dict.get("speaker") or [""])[0] + if isinstance(speaker_for_dialect, str) and speaker_for_dialect.strip(): + spk_is_dialect = getattr(talker_config, "spk_is_dialect", None) or {} + dialect = spk_is_dialect.get(speaker_for_dialect.lower()) + if isinstance(dialect, str) and dialect: + language_id = talker_config.codec_language_id.get(dialect) + if language_id is None: + codec_prefill_list = [ + [ + talker_config.codec_nothink_id, + talker_config.codec_think_bos_id, + talker_config.codec_think_eos_id, + ] + ] + else: + codec_prefill_list = [ + [ + talker_config.codec_think_id, + talker_config.codec_think_bos_id, + int(language_id), + talker_config.codec_think_eos_id, + ] + ] + + codec_input_0 = codec_embed(torch.tensor(codec_prefill_list, device=input_ids.device, dtype=torch.long)) + codec_input_1 = codec_embed( + torch.tensor([[talker_config.codec_pad_id, talker_config.codec_bos_id]], device=input_ids.device) + ) + + # Speaker embedding/token (task-dependent) + speaker_embed: torch.Tensor | None = None + ref_code_len: int | None = None + ref_code_prompt: torch.Tensor | None = None + + def _as_singleton(x: object) -> object: + if isinstance(x, list): + return x[0] if x else None + return x + + def _to_long_tensor(x: object, *, device: torch.device) -> torch.Tensor | None: + x = _as_singleton(x) + if x is None: + return None + if isinstance(x, torch.Tensor): + t = x + elif isinstance(x, np.ndarray): + t = torch.from_numpy(x) + elif isinstance(x, list) and x and all(isinstance(v, (int, np.integer)) for v in x): + t = torch.tensor(x, dtype=torch.long) + else: + return None + if t.ndim == 1: + t = t.unsqueeze(0) + return t.to(device=device, dtype=torch.long) + + def _normalize_voice_clone_prompt(raw: object) -> dict[str, object] | None: + raw = _as_singleton(raw) + if raw is None: + return None + if isinstance(raw, dict): + return raw + # Some callers may pass list[dict] directly. + if isinstance(raw, list) and raw and isinstance(raw[0], dict): + return raw[0] + return None + + if task_type == "Base": + # Base supports voice clone prompt with in-context mode. + xvec_only = bool((info_dict.get("x_vector_only_mode") or [False])[0]) + in_context_mode = not xvec_only + voice_clone_prompt = _normalize_voice_clone_prompt(info_dict.get("voice_clone_prompt")) + ref_audio_wav: np.ndarray | None = None + ref_audio_sr: int | None = None + + def _get_ref_audio() -> tuple[np.ndarray, int]: + nonlocal ref_audio_wav, ref_audio_sr + if ref_audio_wav is None or ref_audio_sr is None: + normalized_ref_audio = info_dict.pop(NORMALIZED_REF_AUDIO_KEY, None) + if ( + isinstance(normalized_ref_audio, tuple) + and len(normalized_ref_audio) == 2 + and isinstance(normalized_ref_audio[0], np.ndarray) + ): + ref_audio_wav = normalized_ref_audio[0] + ref_audio_sr = int(normalized_ref_audio[1]) + return ref_audio_wav, ref_audio_sr + ref_audio_list = info_dict.get("ref_audio") + if not isinstance(ref_audio_list, list) or not ref_audio_list: + raise ValueError("Base requires `ref_audio`.") + ref_audio_wav, ref_audio_sr = self.normalize_ref_audio(ref_audio_list[0]) + return ref_audio_wav, ref_audio_sr + + # Per-ref-audio artifact cache: usually pre-populated by + # preprocess_batch. Fall back to on-demand key creation for + # paths it skips but that still extract an ECAPA spk_embed + # from ref_audio: x_vector_only_mode, precomputed ref_code + # with no ref_spk_embedding, voice_clone_prompt.icl_mode=False. + ref_audio_cache_key = info_dict.pop(REF_AUDIO_CACHE_KEY, None) + if not (isinstance(ref_audio_cache_key, str) and ref_audio_cache_key): + normalized = info_dict.get(NORMALIZED_REF_AUDIO_KEY) + wav_peek: np.ndarray | None = None + sr_peek: int | None = None + if isinstance(normalized, tuple) and len(normalized) == 2 and isinstance(normalized[0], np.ndarray): + wav_peek, sr_peek = normalized[0], int(normalized[1]) + else: + ref_audio_list = info_dict.get("ref_audio") + if isinstance(ref_audio_list, list) and ref_audio_list: + try: + wav_peek, sr_peek = self.normalize_ref_audio(ref_audio_list[0]) + info_dict[NORMALIZED_REF_AUDIO_KEY] = (wav_peek, int(sr_peek)) + except Exception: + wav_peek = None + if isinstance(wav_peek, np.ndarray) and sr_peek is not None: + ref_audio_cache_key = self.make_ref_audio_cache_key(wav_peek, sr_peek) + cached_artifacts = None + if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key: + cached_artifacts = self.get_ref_audio_artifacts(ref_audio_cache_key) + + # Speaker cache: only for uploaded (named) speakers + _speaker_cache_key = None + if voice_clone_prompt is None and self._speaker_cache is not None: + _speaker_list = info_dict.get("speaker") + if isinstance(_speaker_list, list) and _speaker_list: + _voice_name = str(_speaker_list[0]).lower() + # Per-mode namespace — xvec and icl produce different artifacts + # for the same voice, so they must not share a cache slot. + _mode = "xvec" if xvec_only else "icl" + _voice_created_at = int((info_dict.get("voice_created_at") or [0])[0]) + _speaker_cache_key = self._speaker_cache.make_cache_key( + _voice_name, + model_type=f"qwen3_tts_{_mode}", + created_at=_voice_created_at, + ) + _cached = self._speaker_cache.get(_speaker_cache_key) + if _cached is not None: + ref_code_cached = _cached.get("ref_code") + ref_spk_embed_cached = _cached.get("ref_spk_embedding") + if isinstance(ref_code_cached, torch.Tensor): + ref_code_cached = ref_code_cached.to(device=input_ids.device) + if isinstance(ref_spk_embed_cached, torch.Tensor): + ref_spk_embed_cached = ref_spk_embed_cached.to(device=input_ids.device) + voice_clone_prompt = { + "ref_code": ref_code_cached, + "ref_spk_embedding": ref_spk_embed_cached, + "icl_mode": _cached.get("icl_mode"), + } + _speaker_cache_key = None # hit → don't store again + + # Official implementation may pass `voice_clone_prompt.icl_mode`. + if voice_clone_prompt is not None and "icl_mode" in voice_clone_prompt: + icl_flag = _as_singleton(voice_clone_prompt.get("icl_mode")) + if isinstance(icl_flag, bool): + in_context_mode = icl_flag + xvec_only = not in_context_mode + ref_code = None + if voice_clone_prompt is not None: + ref_code = _as_singleton(voice_clone_prompt.get("ref_code")) + ref_code_t = None + if isinstance(ref_code, torch.Tensor): + ref_code_t = ref_code + elif isinstance(ref_code, np.ndarray): + ref_code_t = torch.from_numpy(ref_code) + if isinstance(ref_code_t, torch.Tensor): + if ref_code_t.ndim == 3: + ref_code_t = ref_code_t[0] + ref_code_t = ref_code_t.to(device=input_ids.device, dtype=torch.long) + ref_code_len = int(ref_code_t.shape[0]) + else: + codes = info_dict.get("codes") + precomputed_ref_code = codes.get(PRECOMPUTED_REF_CODE_KEY) if isinstance(codes, dict) else None + ref_code_t = coerce_ref_code_tensor(precomputed_ref_code, device=input_ids.device) + if isinstance(ref_code_t, torch.Tensor): + ref_code_len = int(ref_code_t.shape[0]) + if ref_code_t is None and in_context_mode and cached_artifacts is not None: + cached_ref_code = coerce_ref_code_tensor(cached_artifacts.get("ref_code"), device=input_ids.device) + if isinstance(cached_ref_code, torch.Tensor): + ref_code_t = cached_ref_code + ref_code_len = int(ref_code_t.shape[0]) + if ref_code_t is None and in_context_mode: + # Compute ref_code from ref_audio if not provided. + wav_np, sr = _get_ref_audio() + ref_code_t = self.encode_ref_audio_to_code(wav_np, sr).to(device=input_ids.device) + ref_code_len = int(ref_code_t.shape[0]) + if isinstance(ref_code_t, torch.Tensor): + ref_code_prompt = ref_code_t + + # Speaker embedding: use prompt embed if provided; otherwise extract from audio. + # NOTE: Do NOT use _as_singleton here — the embedding may be a plain + # float list (from API via msgspec IPC) that _as_singleton would + # destructively unwrap to a single scalar. + spk = None + if voice_clone_prompt is not None: + spk = voice_clone_prompt.get("ref_spk_embedding") + if isinstance(spk, torch.Tensor): + speaker_embed = spk.to(device=input_ids.device, dtype=torch.bfloat16).view(1, 1, -1) + elif isinstance(spk, (list, np.ndarray)): + speaker_embed = torch.tensor(spk, dtype=torch.bfloat16, device=input_ids.device).view(1, 1, -1) + elif cached_artifacts is not None and isinstance(cached_artifacts.get("ref_spk_embedding"), torch.Tensor): + speaker_embed = ( + cached_artifacts["ref_spk_embedding"] + .to(device=input_ids.device, dtype=torch.bfloat16) + .view(1, 1, -1) + ) + else: + wav_np, sr = _get_ref_audio() + speaker_embed = self.extract_speaker_embedding(wav_np, sr).view(1, 1, -1) + + # Cache miss: store extraction result in the speaker cache. + if _speaker_cache_key is not None and speaker_embed is not None and self._speaker_cache is not None: + self._speaker_cache.put( + _speaker_cache_key, + { + "ref_code": ref_code_prompt.detach().cpu() + if isinstance(ref_code_prompt, torch.Tensor) + else None, + "ref_spk_embedding": speaker_embed.detach().cpu().reshape(-1), + "icl_mode": in_context_mode, + }, + ) + + # Cache miss: store extraction result in the per-ref-audio artifact cache. + if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key: + self.put_ref_audio_artifacts( + ref_audio_cache_key, + ref_code=ref_code_prompt if isinstance(ref_code_prompt, torch.Tensor) else None, + ref_spk_embedding=speaker_embed.reshape(-1) if isinstance(speaker_embed, torch.Tensor) else None, + ) + + codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1) + + # Role header (<|im_start|>assistant\n) -> projected text embeds. + role_embed = text_projection(text_embedding(input_ids[:, :3])) + + codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1) + codec_prefix = codec_prefix + codec_input[:, :-1] + talker_prompt = torch.cat((role_embed, codec_prefix), dim=1) + + if in_context_mode: + # Prefer explicit tokenized `ref_ids` if provided (matches official signature). + ref_ids = coerce_token_ids(info_dict.pop(PRECOMPUTED_REF_IDS_KEY, None), device=input_ids.device) + if ref_ids is None: + ref_ids = _to_long_tensor(info_dict.get("ref_ids"), device=input_ids.device) + if ref_ids is None and voice_clone_prompt is not None: + ref_ids = _to_long_tensor( + voice_clone_prompt.get("ref_ids") or voice_clone_prompt.get("ref_id"), + device=input_ids.device, + ) + if ref_ids is None: + ref_text = _as_singleton(info_dict.get("ref_text")) + if isinstance(ref_text, str) and ref_text.strip(): + ref_ids = tok( + build_ref_text(ref_text), + return_tensors="pt", + padding=False, + )["input_ids"].to(device=input_ids.device) + else: + logger.warning("Base ICL: ref_text/ref_ids missing, falling back to x-vector-only mode.") + in_context_mode = False + if in_context_mode: + icl_input_embed, trailing_text_hidden = self._generate_icl_prompt( + text_id=input_ids[:, 3:-5], + ref_id=ref_ids[:, 3:-2], + ref_code=ref_code_t, # type: ignore[arg-type] + tts_pad_embed=tts_pad_embed, + tts_eos_embed=tts_eos_embed, + non_streaming_mode=non_streaming_mode, + ) + talker_prompt = torch.cat([talker_prompt, icl_input_embed], dim=1) + else: + # First text token (+ codec_bos). + if non_streaming_mode: + # Official non-streaming mode: put the full text into the + # prefill prompt and use PAD for decode steps. + text_all = text_projection(text_embedding(input_ids[:, 3:-5])) + text_all = torch.cat([text_all, tts_eos_embed], dim=1) + pad_ids = torch.full( + (1, int(text_all.shape[1])), + int(talker_config.codec_pad_id), + device=input_ids.device, + dtype=torch.long, + ) + talker_prompt = torch.cat( + [ + talker_prompt, + text_all + codec_embed(pad_ids), + tts_pad_embed + + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)), + ], + dim=1, + ) + trailing_text_hidden = tts_pad_embed + else: + first_text = text_projection(text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:] + talker_prompt = torch.cat([talker_prompt, first_text], dim=1) + trailing_text_hidden = torch.cat( + ( + text_projection(text_embedding(input_ids[:, 4:-5])), + tts_eos_embed, + ), + dim=1, + ) + + elif task_type == "CustomVoice": + _speaker_raw = info_dict.get("speaker") or [""] + speaker = ( + ((_speaker_raw[0] if isinstance(_speaker_raw, (list, tuple)) else _speaker_raw) or "").lower().strip() + ) + if not speaker: + raise ValueError("CustomVoice requires additional_information.speaker.") + spk_id_map = {k.lower(): v for k, v in (getattr(talker_config, "spk_id", None) or {}).items()} + if speaker not in spk_id_map: + raise ValueError(f"Unsupported speaker: {speaker}") + spk_id = spk_id_map[speaker] + # Keep it at least 1D; embedding on a 0-d tensor can return 1D. + spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long) + spk_embed = codec_embed(spk_tensor) + if spk_embed.ndim == 1: + spk_embed = spk_embed.view(1, 1, -1) + elif spk_embed.ndim == 2: + spk_embed = spk_embed.view(1, 1, -1) + speaker_embed = spk_embed + codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1) + + role_embed = text_projection(text_embedding(input_ids[:, :3])) + codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1) + codec_prefix = codec_prefix + codec_input[:, :-1] + talker_prompt = torch.cat((role_embed, codec_prefix), dim=1) + + if non_streaming_mode: + text_all = text_projection(text_embedding(input_ids[:, 3:-5])) + text_all = torch.cat([text_all, tts_eos_embed], dim=1) + pad_ids = torch.full( + (1, int(text_all.shape[1])), + int(talker_config.codec_pad_id), + device=input_ids.device, + dtype=torch.long, + ) + talker_prompt = torch.cat( + [ + talker_prompt, + text_all + codec_embed(pad_ids), + tts_pad_embed + + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)), + ], + dim=1, + ) + trailing_text_hidden = tts_pad_embed + else: + first_text = text_projection(text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:] + talker_prompt = torch.cat([talker_prompt, first_text], dim=1) + trailing_text_hidden = torch.cat( + ( + text_projection(text_embedding(input_ids[:, 4:-5])), + tts_eos_embed, + ), + dim=1, + ) + + elif task_type == "VoiceDesign": + # No known speaker identity; only codec tags + text. + codec_input = torch.cat([codec_input_0, codec_input_1], dim=1) + + role_embed = text_projection(text_embedding(input_ids[:, :3])) + codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1) + codec_prefix = codec_prefix + codec_input[:, :-1] + talker_prompt = torch.cat((role_embed, codec_prefix), dim=1) + + if non_streaming_mode: + text_all = text_projection(text_embedding(input_ids[:, 3:-5])) + text_all = torch.cat([text_all, tts_eos_embed], dim=1) + pad_ids = torch.full( + (1, int(text_all.shape[1])), + int(talker_config.codec_pad_id), + device=input_ids.device, + dtype=torch.long, + ) + talker_prompt = torch.cat( + [ + talker_prompt, + text_all + codec_embed(pad_ids), + tts_pad_embed + + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)), + ], + dim=1, + ) + trailing_text_hidden = tts_pad_embed + else: + first_text = text_projection(text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:] + talker_prompt = torch.cat([talker_prompt, first_text], dim=1) + trailing_text_hidden = torch.cat( + ( + text_projection(text_embedding(input_ids[:, 4:-5])), + tts_eos_embed, + ), + dim=1, + ) + else: + raise ValueError(f"Unsupported task_type={task_type}") + + if instruct_embed is not None: + talker_prompt = torch.cat([instruct_embed, talker_prompt], dim=1) + + return ( + talker_prompt.squeeze(0), # [prompt_len, H] + trailing_text_hidden.squeeze(0), # [T, H] + ref_code_len, + ref_code_prompt.contiguous() if isinstance(ref_code_prompt, torch.Tensor) else None, + ) + + # -------------------- prompt-length estimator -------------------- + + @staticmethod + def estimate_prompt_len_from_additional_information( + additional_information: dict[str, Any] | None, + *, + task_type: str, + tokenize_prompt: Callable[[str], list[int]], + codec_language_id: Mapping[str, int] | None, + spk_is_dialect: Mapping[str, object] | None, + estimate_ref_code_len: Callable[[object], int | None] | None = None, + ) -> int: + """Compute Stage-0 placeholder prompt length (length-only mirror of + :meth:`build_prompt_embeds`). + + It must match the model-side ``inputs_embeds`` length to avoid extra + padding and quality drop. + """ + + def _first(x: object, default: object) -> object: + if isinstance(x, list): + return x[0] if x else default + return x if x is not None else default + + info: dict[str, Any] = additional_information or {} + text = _first(info.get("text"), "") + language = _first(info.get("language"), "Auto") + speaker = _first(info.get("speaker"), "").lower().strip() + instruct = _first(info.get("instruct"), "") + non_streaming_mode_raw = _first(info.get("non_streaming_mode"), None) + + if isinstance(non_streaming_mode_raw, bool): + non_streaming_mode = non_streaming_mode_raw + else: + # Official defaults: CustomVoice/VoiceDesign -> non_streaming_mode=True; Base -> False. + non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign") + + if not isinstance(text, str): + text = "" + if not isinstance(instruct, str): + instruct = "" + if not isinstance(language, str): + language = "Auto" + + instruct_len = 0 + if instruct.strip(): + instruct_len = len(tokenize_prompt(build_instruct_text(instruct))) + + # ---- codec prefix portion (matches build_prompt_embeds) ---- + language_id = None + if language.lower() != "auto" and codec_language_id: + language_id = codec_language_id.get(language.lower()) + if ( + language_id is None + and codec_language_id + and spk_is_dialect + and isinstance(language, str) + and language.lower() in ("chinese", "auto") + and isinstance(speaker, str) + and speaker.strip() + ): + dialect = spk_is_dialect.get(speaker.lower()) + if isinstance(dialect, str) and dialect: + language_id = codec_language_id.get(dialect) + prefill_len = 3 if language_id is None else 4 + + speaker_len = 1 if task_type in ("CustomVoice", "Base") else 0 + codec_input_len = prefill_len + speaker_len + 2 # + [codec_pad, codec_bos] + codec_prefix_len = codec_input_len - 1 # codec_input[:-1] + tts_bos + + # Role header: input_ids[:, :3] in model. + role_len = 3 + prompt_len = instruct_len + role_len + codec_prefix_len + + # ---- text conditioning portion (matches build_prompt_embeds) ---- + assistant_len = len(tokenize_prompt(build_assistant_text(text))) + if assistant_len < 8: + raise ValueError(f"Unexpected assistant prompt length: {assistant_len}") + + if task_type in ("CustomVoice", "VoiceDesign"): + if non_streaming_mode: + # model: full text ids (input_ids[:, 3:-5]) + eos + codec_bos step + prompt_len += assistant_len - 6 + else: + # model: only first text token in prefill + prompt_len += 1 + + if task_type == "Base": + xvec_only = bool(_first(info.get("x_vector_only_mode"), False)) + in_context_mode = not xvec_only + + voice_clone_prompt = _first(info.get("voice_clone_prompt"), None) + if isinstance(voice_clone_prompt, dict): + icl_flag = _first(voice_clone_prompt.get("icl_mode"), None) + if isinstance(icl_flag, bool): + in_context_mode = icl_flag + + if in_context_mode: + ref_code = None + if isinstance(voice_clone_prompt, dict): + # Do NOT apply `_first(...)` here. `_first` unwraps the singleton-list + # batching convention used at the top level of `additional_information`; + # values inside `voice_clone_prompt` are per-item payloads. `ref_code` + # is a `[num_frames, num_codebooks]` 2D list, so `_first` would strip + # its outer dim and report `num_codebooks` (~8) instead of `num_frames` + # (~hundreds), silently truncating the prefill embeddings downstream. + ref_code = voice_clone_prompt.get("ref_code") + + ref_code_len: int | None = None + if isinstance(ref_code, list): + if ref_code and isinstance(ref_code[0], list): + ref_code_len = len(ref_code) + elif ref_code: + ref_code_len = len(ref_code) + elif hasattr(ref_code, "shape"): + try: + shape = getattr(ref_code, "shape") + if shape and len(shape) >= 1: + ref_code_len = int(shape[0]) + except Exception: + ref_code_len = None + + if ref_code_len is None and estimate_ref_code_len is not None: + ref_code_len = estimate_ref_code_len(info.get("ref_audio")) + if ref_code_len is None: + raise ValueError( + "Base in-context voice cloning requires either `voice_clone_prompt.ref_code` " + "or a readable `ref_audio` that can be mapped to a codec frame length." + ) + + codec_lens = 1 + int(ref_code_len) # codec_bos + ref_code + if non_streaming_mode: + # Mirrors _generate_icl_prompt(non_streaming_mode=True): + # text_embed = ref_ids + text_ids + eos. + ref_ids = _first(info.get("ref_ids"), None) + if isinstance(voice_clone_prompt, dict) and ref_ids is None: + # Same reasoning as ref_code above: values inside `voice_clone_prompt` + # are per-item payloads, not singleton-wrapped batches. Applying + # `_first` to a list-of-int `ref_ids` collapses it to an int and the + # branches below fall through to `ref_ids_len = 0`. + ref_ids = voice_clone_prompt.get("ref_ids") or voice_clone_prompt.get("ref_id") + + if ref_ids is None: + ref_text = _first(info.get("ref_text"), "") + if not isinstance(ref_text, str) or not ref_text.strip(): + raise ValueError( + "Base in-context non-streaming requires `ref_text` or tokenized `ref_ids`." + ) + ref_ids_len = len(tokenize_prompt(build_ref_text(ref_text))) + elif hasattr(ref_ids, "shape"): + shape = getattr(ref_ids, "shape", None) + ref_ids_len = int(shape[-1]) if shape else 0 + elif isinstance(ref_ids, list): + ref_ids_len = len(ref_ids) + else: + ref_ids_len = 0 + + # model uses ref_ids[:, 3:-2] (strip 5 tokens) and text_id=input_ids[:, 3:-5] (strip 8). + ref_id_len = max(0, int(ref_ids_len) - 5) + text_id_len = max(0, int(assistant_len) - 8) + text_embed_len = ref_id_len + text_id_len + 1 # + eos + prompt_len += text_embed_len + codec_lens + else: + # Mirrors _generate_icl_prompt(non_streaming_mode=False): + # aligned to codec_lens. + prompt_len += codec_lens + else: + # Base without ICL behaves like CustomVoice. + if non_streaming_mode: + prompt_len += assistant_len - 6 + else: + prompt_len += 1 + + return max(2, int(prompt_len)) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index bd5ec6108..ca1c4b395 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -1,21 +1,14 @@ from __future__ import annotations -import base64 import copy -import hashlib -import io -from collections import OrderedDict -from collections.abc import Callable, Iterable, Mapping -from functools import lru_cache +from collections.abc import Iterable, Mapping from typing import Any -from urllib.parse import urlparse import numpy as np -import soundfile as sf import torch import torch.nn as nn import torch.nn.functional as F -from transformers import AutoFeatureExtractor, AutoTokenizer +from transformers import AutoFeatureExtractor from transformers.activations import ACT2FN from vllm.config import VllmConfig from vllm.distributed import get_pp_group @@ -30,10 +23,10 @@ from vllm.sequence import IntermediateTensors from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.utils.audio import mel_filter_bank from vllm_omni.utils.speaker_cache import get_speaker_cache from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig +from .prompt_embeds_builder import Qwen3TTSPromptEmbedsBuilder from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder @@ -41,11 +34,6 @@ from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2E logger = init_logger(__name__) _TRAILING_TEXT_COMPACT_MIN_FRAMES = 64 -_PRECOMPUTED_REF_CODE_KEY = "precomputed_ref" -_NORMALIZED_REF_AUDIO_KEY = "_qwen3_tts_normalized_ref_audio" -_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key" -_PRECOMPUTED_TEXT_IDS_KEY = "_qwen3_tts_text_ids" -_PRECOMPUTED_REF_IDS_KEY = "_qwen3_tts_ref_ids" # --------------------------------------------------------------------------- @@ -252,69 +240,6 @@ class Qwen3TTSSpeakerEncoder(torch.nn.Module): return hidden_states.squeeze(-1) -# ---- Audio utilities ---- - - -def _dynamic_range_compression(x, c=1, clip_val=1e-5): - return torch.log(torch.clamp(x, min=clip_val) * c) - - -@lru_cache(maxsize=8) -def _cached_mel_filter_bank(sampling_rate: int, n_fft: int, n_mels: int, fmin: int, fmax: int | None) -> torch.Tensor: - return mel_filter_bank(sr=sampling_rate, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax) - - -@lru_cache(maxsize=8) -def _cached_hann_window(win_size: int) -> torch.Tensor: - return torch.hann_window(win_size) - - -def mel_spectrogram( - y: torch.Tensor, - n_fft: int, - num_mels: int, - sampling_rate: int, - hop_size: int, - win_size: int, - fmin: int, - fmax: int | None = None, - center: bool = False, - mel_basis: torch.Tensor | None = None, - hann_window: torch.Tensor | None = None, -) -> torch.Tensor: - """Calculate mel spectrogram of an input signal using torchaudio mel filterbank and torch STFT.""" - if torch.min(y) < -1.0: - logger.warning("Min value of input waveform signal is %s", torch.min(y)) - if torch.max(y) > 1.0: - logger.warning("Max value of input waveform signal is %s", torch.max(y)) - device = y.device - if mel_basis is None: - mel_basis = _cached_mel_filter_bank(sampling_rate, n_fft, num_mels, fmin, fmax).to(device) - elif mel_basis.device != device: - mel_basis = mel_basis.to(device) - if hann_window is None: - hann_window = _cached_hann_window(win_size).to(device) - elif hann_window.device != device: - hann_window = hann_window.to(device) - padding = (n_fft - hop_size) // 2 - y = torch.nn.functional.pad(y.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1) - spec = torch.stft( - y, - n_fft, - hop_length=hop_size, - win_length=win_size, - window=hann_window, - center=center, - pad_mode="reflect", - normalized=False, - onesided=True, - return_complex=True, - ) - spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9) - mel_spec = torch.matmul(mel_basis, spec) - return _dynamic_range_compression(mel_spec) - - # --------------------------------------------------------------------------- # Main AR Talker model # --------------------------------------------------------------------------- @@ -453,12 +378,23 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): self.gpu_resident_buffer_keys: set[tuple[str, str]] = { ("codes", "audio"), ("hidden_states", "last"), - ("embed", "tts_pad"), ("hidden_states", "trailing_text"), } - # Tokenizer for prompt building. - self._tokenizer = None + # ``text_proj(text_emb(tts_pad_token_id))`` is request-independent — + # it depends only on frozen ``text_embedding`` / ``text_projection`` + # weights, so we precompute it once in :meth:`_init_runtime_buffers` + # (called from :meth:`load_weights`) and reuse the same buffer at + # every prefill and decode step instead of round-tripping it through + # the per-request ``info_dict``. Declared here as zeros so the + # attribute exists under ``load_format: dummy`` (which bypasses + # ``load_weights`` entirely and leaves the value uninitialized). + model_dtype = getattr(vllm_config.model_config, "dtype", torch.bfloat16) + self.register_buffer( + "_tts_pad_embed", + torch.zeros(1, int(self.talker_config.hidden_size), dtype=model_dtype), + persistent=False, + ) tokenizer_config = Qwen3TTSTokenizerV2Config.from_pretrained( self.model_path, @@ -478,36 +414,39 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): ) self._speaker_cache = get_speaker_cache() - self._ref_audio_artifact_cache_max_entries = 256 - self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict() raw_subtalker_sampling = getattr(vllm_config.model_config, "subtalker_sampling_params", None) self._subtalker_sampling_params: dict[str, Any] = ( dict(raw_subtalker_sampling) if isinstance(raw_subtalker_sampling, Mapping) else {} ) self._stacked_codec_embed: torch.Tensor | None = None - # Bounded LRU: caller-supplied orig_sr can otherwise grow this without limit. - self._resampler_cache: OrderedDict[tuple[int, int], AudioResampler] = OrderedDict() - self._resampler_cache_max = 16 - self._tts_pad_embed_cache: torch.Tensor | None = None + + # Stand-alone wrapper around the embedding layers, encoders and + # tokenizers required to assemble a talker prefill prompt. Owns the + # text + speech tokenizers (loaded lazily on first use) and the + # per-ref-audio / resampler caches, so they don't leak onto the + # talker class. Other talker variants can construct their own builder + # with the same set of dependencies and reuse ``build_prompt_embeds`` + # verbatim. + self._prompt_builder = Qwen3TTSPromptEmbedsBuilder( + config=self.config, + talker_config=self.talker_config, + model_path=self.model_path, + text_embedding=self.text_embedding, + text_projection=self.text_projection, + codec_embed=self.embed_input_ids, + residual_code_embeddings=lambda: self.code_predictor.get_input_embeddings(), + speaker_encoder=self.speaker_encoder, + tts_pad_embed=self._tts_pad_embed, + encode_ref_audio_batch=self._encode_ref_audio_batch, + speaker_cache=self._speaker_cache, + ) # -------------------- vLLM required hooks -------------------- def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: return self.model.embed_input_ids(input_ids) - def _get_tts_pad_embed_cached(self, device: torch.device) -> torch.Tensor: - """Return a device-pinned tts_pad embedding, rebuilding once per device.""" - cached = self._tts_pad_embed_cache - if cached is not None and cached.device == device: - return cached - pad_id = torch.tensor([[self.config.tts_pad_token_id]], device=device, dtype=torch.long) - with torch.no_grad(): - projected = self.text_projection(self.text_embedding(pad_id)) - cached = projected.reshape(1, -1).to(dtype=torch.bfloat16).detach() - self._tts_pad_embed_cache = cached - return cached - def forward( self, input_ids: torch.Tensor, @@ -654,32 +593,30 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): else: codec_streaming = task_type == "Base" + # ``tts_pad_embed`` is a request-independent constant — see + # :meth:`_init_runtime_buffers`. Materialize once on the right + # device/dtype and reuse for both the prefill placeholder padding + # and the decode text-step fallback below. + tts_pad_embed = self._tts_pad_embed.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1) + if is_prefill: # Prefill (prompt embeddings) prompt_embeds_cpu = embed.get("prefill") - tts_pad_embed_cpu = embed.get("tts_pad") - tts_pad_embed = None - if isinstance(tts_pad_embed_cpu, torch.Tensor) and tts_pad_embed_cpu.numel() > 0: - tts_pad_embed = tts_pad_embed_cpu.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1) # First prefill round: prompt_embeds_cpu is not yet populated. # Subsequent prefill rounds (multi-chunk): prompt_embeds_cpu is a Tensor stored by the first round. is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2 if is_first_prefill: - full_prompt_embeds, tailing_text_hidden, tts_pad_embed, ref_code_len, ref_code = ( - self._build_prompt_embeds(task_type=task_type, info_dict=info_dict) + full_prompt_embeds, tailing_text_hidden, ref_code_len, ref_code = ( + self._prompt_builder.build_prompt_embeds(task_type=task_type, info_dict=info_dict) ) # Store full prompt embeddings on CPU (large, prefill-only). - # tailing_text_hidden and tts_pad_embed stay on GPU (gpu_resident_buffer_keys). + # tailing_text_hidden stays on GPU (gpu_resident_buffer_keys). prompt_embeds_cpu = full_prompt_embeds.detach().to("cpu").contiguous() info_update: OmniPayload = { - "embed": { - "prefill": prompt_embeds_cpu, - "tts_pad": tts_pad_embed.detach(), - }, + "embed": {"prefill": prompt_embeds_cpu}, "hidden_states": {"trailing_text": tailing_text_hidden.detach()}, "meta": { - "talker_prefill_offset": 0, "talker_text_offset": 0, "codec_streaming": codec_streaming, }, @@ -688,41 +625,27 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): info_update.setdefault("codes", {})["ref"] = ref_code.detach().to("cpu").contiguous() if ref_code_len is not None: info_update["meta"]["ref_code_len"] = int(ref_code_len) - # Always return a span_len slice; if the scheduled placeholder is longer, pad with tts_pad_embed. - # This preserves placeholder/embedding alignment. - offset = int(info_dict.get("_omni_num_computed_tokens", 0) or 0) - if offset < 0: - offset = 0 - s = max(0, min(offset, int(prompt_embeds_cpu.shape[0]))) - e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0]))) - take = prompt_embeds_cpu[s:e] - if int(take.shape[0]) < span_len: - pad_n = int(span_len - int(take.shape[0])) - pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1) - take = torch.cat([take, pad_rows], dim=0) - prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16) - info_update["meta"]["talker_prefill_offset"] = int(offset + span_len) + # First prefill: source the slice offset from `_omni_num_computed_tokens` + # so cache-recovery (prefill replay at a later offset) lands on the right + # slice of the stored embeddings. Subsequent chunks below advance from + # `talker_prefill_offset` written here. + offset = max(0, int(info_dict.get("_omni_num_computed_tokens", 0) or 0)) else: # Subsequent prefill chunk: slice from stored embeddings at running offset. - if tts_pad_embed is None: - raise RuntimeError("Missing `tts_pad_embed` in additional_information; prefill must initialize it.") - offset = int(meta.get("talker_prefill_offset", 0) or 0) - if offset < 0: - offset = 0 - s = max(0, min(offset, int(prompt_embeds_cpu.shape[0]))) - e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0]))) - take = prompt_embeds_cpu[s:e] - if int(take.shape[0]) < span_len: - pad_n = int(span_len - int(take.shape[0])) - pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1) - take = torch.cat([take, pad_rows], dim=0) - prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16) - info_update = { - "meta": { - "talker_prefill_offset": int(offset + span_len), - "codec_streaming": codec_streaming, - } - } + offset = max(0, int(meta.get("talker_prefill_offset", 0) or 0)) + info_update = {"meta": {"codec_streaming": codec_streaming}} + + # Always return a span_len slice; if the scheduled placeholder is longer than what + # the prompt actually fills, pad with tts_pad_embed (preserves placeholder/embedding alignment). + s = max(0, min(offset, int(prompt_embeds_cpu.shape[0]))) + e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0]))) + take = prompt_embeds_cpu[s:e] + if int(take.shape[0]) < span_len: + pad_n = int(span_len - int(take.shape[0])) + pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1) + take = torch.cat([take, pad_rows], dim=0) + prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16) + info_update["meta"]["talker_prefill_offset"] = int(offset + span_len) # When inputs_embeds is set, token ids are ignored by the model but must stay in-vocab for vLLM bookkeeping. input_ids_out = input_ids.clone() @@ -736,13 +659,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): info_update.setdefault("codes", {})["audio"] = zeros return input_ids_out, prompt_embeds, info_update - # Decode: span_len == 1. Buffers are GPU-resident so .to() is a no-op. - tts_pad_embed_buf = embed.get("tts_pad") - if isinstance(tts_pad_embed_buf, torch.Tensor): - tts_pad_embed = tts_pad_embed_buf.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1) - else: - # Defensive: rebuild from text_embedding when prefill state was evicted. - tts_pad_embed = self._get_tts_pad_embed_cached(input_ids.device) + # Decode: span_len == 1 + # Pop one text-step vector from tailing_text_hidden queue. + # ``tts_pad_embed`` was materialized above from :attr:`_tts_pad_embed` + # (request-independent buffer) — no per-request fetch needed. tail = hs.get("trailing_text") text_offset = max(0, int(meta.get("talker_text_offset", 0) or 0)) @@ -820,6 +740,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): device = input_ids_flat.device dtype = torch.bfloat16 + # Request-independent constant (see :meth:`_init_runtime_buffers`) — + # compute once for the batch instead of fetching it per request from + # ``info_dict["embed"]["tts_pad"]``. + tts_pad_embed = self._tts_pad_embed.to(device=device, dtype=dtype).reshape(1, -1) past_hidden_list: list[torch.Tensor] = [] text_step_list: list[torch.Tensor] = [] updates: list[dict[str, Any]] = [] @@ -833,7 +757,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): info_dict = merged payload: OmniPayload = info_dict - embed = payload.get("embed", {}) hs = payload.get("hidden_states", {}) meta = payload.get("meta", {}) @@ -852,11 +775,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): else: codec_streaming = task_type == "Base" - tts_pad_embed_buf = embed.get("tts_pad") - if not isinstance(tts_pad_embed_buf, torch.Tensor): - raise RuntimeError("Missing `tts_pad_embed` in additional_information; prefill must run first.") - tts_pad_embed = tts_pad_embed_buf.to(device=device, dtype=dtype).reshape(1, -1) - tail = hs.get("trailing_text") text_offset = max(0, int(meta.get("talker_text_offset", 0) or 0)) trailing_text_update = None @@ -920,504 +838,20 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): last = hidden_states[-1, :].detach() return {"hidden_states": {"last": last}} - # -------------------- prompt construction helpers -------------------- - - def _get_tokenizer(self): - if self._tokenizer is None: - import transformers - - kwargs = dict(trust_remote_code=True, use_fast=True) - if transformers.__version__ < "5": - kwargs["fix_mistral_regex"] = True - self._tokenizer = AutoTokenizer.from_pretrained(self.model_path, **kwargs) - self._tokenizer.padding_side = "left" - return self._tokenizer - - @staticmethod - def _build_assistant_text(text: str) -> str: - return f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" - - @staticmethod - def _build_ref_text(text: str) -> str: - return f"<|im_start|>assistant\n{text}<|im_end|>\n" - - @staticmethod - def _build_instruct_text(instruct: str) -> str: - return f"<|im_start|>user\n{instruct}<|im_end|>\n" - - @staticmethod - def estimate_prompt_len_from_additional_information( - additional_information: dict[str, Any] | None, - *, - task_type: str, - tokenize_prompt: Callable[[str], list[int]], - codec_language_id: Mapping[str, int] | None, - spk_is_dialect: Mapping[str, object] | None, - estimate_ref_code_len: Callable[[object], int | None] | None = None, - ) -> int: - """Compute Stage-0 placeholder prompt length (length-only mirror of `_build_prompt_embeds()`). - It must match the model-side `inputs_embeds` length to avoid extra padding and quality drop.""" - - def _first(x: object, default: object) -> object: - if isinstance(x, list): - return x[0] if x else default - return x if x is not None else default - - info: dict[str, Any] = additional_information or {} - text = _first(info.get("text"), "") - language = _first(info.get("language"), "Auto") - speaker = _first(info.get("speaker"), "").lower().strip() - instruct = _first(info.get("instruct"), "") - non_streaming_mode_raw = _first(info.get("non_streaming_mode"), None) - - if isinstance(non_streaming_mode_raw, bool): - non_streaming_mode = non_streaming_mode_raw - else: - # Official defaults: CustomVoice/VoiceDesign -> non_streaming_mode=True; Base -> False. - non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign") - - if not isinstance(text, str): - text = "" - if not isinstance(instruct, str): - instruct = "" - if not isinstance(language, str): - language = "Auto" - - instruct_len = 0 - if instruct.strip(): - instruct_text = Qwen3TTSTalkerForConditionalGeneration._build_instruct_text(instruct) - instruct_len = len(tokenize_prompt(instruct_text)) - - # ---- codec prefix portion (matches _build_prompt_embeds) ---- - language_id = None - if language.lower() != "auto" and codec_language_id: - language_id = codec_language_id.get(language.lower()) - if ( - language_id is None - and codec_language_id - and spk_is_dialect - and isinstance(language, str) - and language.lower() in ("chinese", "auto") - and isinstance(speaker, str) - and speaker.strip() - ): - dialect = spk_is_dialect.get(speaker.lower()) - if isinstance(dialect, str) and dialect: - language_id = codec_language_id.get(dialect) - prefill_len = 3 if language_id is None else 4 - - speaker_len = 1 if task_type in ("CustomVoice", "Base") else 0 - codec_input_len = prefill_len + speaker_len + 2 # + [codec_pad, codec_bos] - codec_prefix_len = codec_input_len - 1 # codec_input[:-1] + tts_bos - - # Role header: input_ids[:, :3] in model. - role_len = 3 - prompt_len = instruct_len + role_len + codec_prefix_len - - # ---- text conditioning portion (matches _build_prompt_embeds) ---- - assistant_text = Qwen3TTSTalkerForConditionalGeneration._build_assistant_text(text) - assistant_len = len(tokenize_prompt(assistant_text)) - if assistant_len < 8: - raise ValueError(f"Unexpected assistant prompt length: {assistant_len}") - - if task_type in ("CustomVoice", "VoiceDesign"): - if non_streaming_mode: - # model: full text ids (input_ids[:, 3:-5]) + eos + codec_bos step - prompt_len += assistant_len - 6 - else: - # model: only first text token in prefill - prompt_len += 1 - - if task_type == "Base": - xvec_only = bool(_first(info.get("x_vector_only_mode"), False)) - in_context_mode = not xvec_only - - voice_clone_prompt = _first(info.get("voice_clone_prompt"), None) - if isinstance(voice_clone_prompt, dict): - icl_flag = _first(voice_clone_prompt.get("icl_mode"), None) - if isinstance(icl_flag, bool): - in_context_mode = icl_flag - - if in_context_mode: - ref_code = None - if isinstance(voice_clone_prompt, dict): - # Do NOT apply `_first(...)` here. `_first` unwraps the singleton-list - # batching convention used at the top level of `additional_information`; - # values inside `voice_clone_prompt` are per-item payloads. `ref_code` - # is a `[num_frames, num_codebooks]` 2D list, so `_first` would strip - # its outer dim and report `num_codebooks` (~8) instead of `num_frames` - # (~hundreds), silently truncating the prefill embeddings downstream. - ref_code = voice_clone_prompt.get("ref_code") - - ref_code_len: int | None = None - if isinstance(ref_code, list): - if ref_code and isinstance(ref_code[0], list): - ref_code_len = len(ref_code) - elif ref_code: - ref_code_len = len(ref_code) - elif hasattr(ref_code, "shape"): - try: - shape = getattr(ref_code, "shape") - if shape and len(shape) >= 1: - ref_code_len = int(shape[0]) - except Exception: - ref_code_len = None - - if ref_code_len is None and estimate_ref_code_len is not None: - ref_code_len = estimate_ref_code_len(info.get("ref_audio")) - if ref_code_len is None: - raise ValueError( - "Base in-context voice cloning requires either `voice_clone_prompt.ref_code` " - "or a readable `ref_audio` that can be mapped to a codec frame length." - ) - - codec_lens = 1 + int(ref_code_len) # codec_bos + ref_code - if non_streaming_mode: - # _generate_icl_prompt(non_streaming_mode=True): - # text_embed = ref_ids + text_ids + eos. - ref_ids = _first(info.get("ref_ids"), None) - if isinstance(voice_clone_prompt, dict) and ref_ids is None: - # Same reasoning as ref_code above: values inside `voice_clone_prompt` - # are per-item payloads, not singleton-wrapped batches. Applying - # `_first` to a list-of-int `ref_ids` collapses it to an int and the - # branches below fall through to `ref_ids_len = 0`. - ref_ids = voice_clone_prompt.get("ref_ids") or voice_clone_prompt.get("ref_id") - - if ref_ids is None: - ref_text = _first(info.get("ref_text"), "") - if not isinstance(ref_text, str) or not ref_text.strip(): - raise ValueError( - "Base in-context non-streaming requires `ref_text` or tokenized `ref_ids`." - ) - ref_text_ids = tokenize_prompt(Qwen3TTSTalkerForConditionalGeneration._build_ref_text(ref_text)) - ref_ids_len = len(ref_text_ids) - elif hasattr(ref_ids, "shape"): - shape = getattr(ref_ids, "shape", None) - ref_ids_len = int(shape[-1]) if shape else 0 - elif isinstance(ref_ids, list): - ref_ids_len = len(ref_ids) - else: - ref_ids_len = 0 - - # model uses ref_ids[:, 3:-2] (strip 5 tokens) and text_id=input_ids[:, 3:-5] (strip 8). - ref_id_len = max(0, int(ref_ids_len) - 5) - text_id_len = max(0, int(assistant_len) - 8) - text_embed_len = ref_id_len + text_id_len + 1 # + eos - prompt_len += text_embed_len + codec_lens - else: - # _generate_icl_prompt(non_streaming_mode=False): aligned to codec_lens. - prompt_len += codec_lens - else: - # Base without ICL behaves like CustomVoice. - if non_streaming_mode: - prompt_len += assistant_len - 6 - else: - prompt_len += 1 - - return max(2, int(prompt_len)) - - def _is_probably_base64(self, s: str) -> bool: - if s.startswith("data:audio"): - return True - if ("/" not in s and "\\" not in s) and len(s) > 256: - return True - return False - - def _is_url(self, s: str) -> bool: - try: - u = urlparse(s) - if u.scheme in ("http", "https"): - return bool(u.netloc) - return u.scheme == "file" - except Exception: - return False - - def _decode_base64_to_wav_bytes(self, b64: str) -> bytes: - if "," in b64 and b64.strip().startswith("data:"): - b64 = b64.split(",", 1)[1] - return base64.b64decode(b64) - - def _load_audio_to_np(self, x: str) -> tuple[np.ndarray, int]: - """Load audio from local path, URL, or base64 data URI. - - Uses upstream vLLM's MediaConnector for http(s) URLs and ``file:`` - URIs, with unrestricted local access (offline inference is trusted). - """ - from vllm.multimodal.media.audio import load_audio - - if self._is_url(x): - from vllm.multimodal.media import MediaConnector - - connector = MediaConnector(allowed_local_media_path="/") - audio, sr = connector.fetch_audio(x) - elif self._is_probably_base64(x): - wav_bytes = self._decode_base64_to_wav_bytes(x) - with io.BytesIO(wav_bytes) as f: - audio, sr = sf.read(f, dtype="float32", always_2d=False) - else: - audio, sr = load_audio(x, sr=None, mono=True) - - if isinstance(audio, np.ndarray) and audio.ndim > 1: - audio = np.mean(audio, axis=-1) - - return np.asarray(audio, dtype=np.float32), int(sr) - - def _normalize_ref_audio(self, ref_audio: object) -> tuple[np.ndarray, int]: - # NOTE: additional_information may serialize (wav, sr) into (nested) lists across processes; be tolerant. - if isinstance(ref_audio, str): - return self._load_audio_to_np(ref_audio) - - def _is_sr(x: object) -> bool: - try: - v = int(x) # type: ignore[arg-type] - except Exception: - return False - return 1_000 <= v <= 200_000 - - def _is_number_sequence(xs: list[object]) -> bool: - if not xs: - return False - for v in xs[:8]: - if not isinstance(v, (int, float, np.number)): - return False - return True - - wav_candidates: list[object] = [] - sr_candidates: list[int] = [] - - def _summarize(obj: object, depth: int = 0) -> str: - if depth > 2: - if isinstance(obj, (int, np.integer)): - return f"int({int(obj)})" - return type(obj).__name__ - if obj is None: - return "None" - if isinstance(obj, str): - if len(obj) <= 16: - return f"str({obj!r})" - return f"str(len={len(obj)})" - if isinstance(obj, (int, float, np.number)): - return f"{type(obj).__name__}({obj})" - if isinstance(obj, np.ndarray): - return f"ndarray(shape={obj.shape}, dtype={obj.dtype})" - if isinstance(obj, torch.Tensor): - return f"Tensor(shape={tuple(obj.shape)}, dtype={obj.dtype}, device={obj.device})" - if isinstance(obj, dict): - keys = list(obj.keys()) - return f"dict(keys={keys[:8]})" - if isinstance(obj, (tuple, list)): - items = list(obj) - head = ", ".join(_summarize(x, depth + 1) for x in items[:3]) - return f"{type(obj).__name__}(len={len(items)}; head=[{head}])" - return f"{type(obj).__name__}" - - def _scan(obj: object, depth: int = 0) -> None: - if depth > 4: - return - if obj is None: - return - if _is_sr(obj): - sr_candidates.append(int(obj)) # type: ignore[arg-type] - return - if isinstance(obj, np.ndarray) and obj.size > 0: - wav_candidates.append(obj) - return - if isinstance(obj, torch.Tensor) and obj.numel() > 0: - wav_candidates.append(obj) - return - if isinstance(obj, dict): - # Inlined ndarray/tensor payloads from the input processor. - if obj.get("__ndarray__") and "data" in obj and "dtype" in obj and "shape" in obj: - try: - data = obj["data"] - dtype = obj["dtype"] - shape = obj["shape"] - if isinstance(data, (bytes, bytearray, memoryview)): - arr = np.frombuffer(data, dtype=dtype).reshape(shape) - if arr.size > 0: - wav_candidates.append(arr) - return - except Exception: - pass - if obj.get("__tensor__") and "data" in obj and "dtype" in obj and "shape" in obj: - try: - data = obj["data"] - dtype = obj["dtype"] - shape = obj["shape"] - if isinstance(data, (bytes, bytearray, memoryview)): - # Stored as raw CPU bytes; interpret as numpy for audio. - np_dtype = np.dtype(dtype) - arr = np.frombuffer(data, dtype=np_dtype).reshape(shape) - if arr.size > 0: - wav_candidates.append(arr) - return - except Exception: - pass - wav_obj = obj.get("array") or obj.get("wav") or obj.get("audio") - sr_obj = obj.get("sampling_rate") or obj.get("sr") or obj.get("sample_rate") - if wav_obj is not None: - _scan(wav_obj, depth + 1) - if sr_obj is not None: - _scan(sr_obj, depth + 1) - return - if isinstance(obj, (tuple, list)): - obj_list = list(obj) - # Unwrap singleton nesting ([[wav, sr]]). - while isinstance(obj_list, list) and len(obj_list) == 1: - inner = obj_list[0] - if isinstance(inner, np.ndarray) and inner.size > 0: - wav_candidates.append(inner) - return - if isinstance(inner, torch.Tensor) and inner.numel() > 0: - wav_candidates.append(inner) - return - if isinstance(inner, dict): - _scan(inner, depth + 1) - return - if isinstance(inner, (tuple, list)): - obj_list = list(inner) # type: ignore[list-item] - continue - break - - # If the *unwrapped* list is a long list of numbers, treat it as waveform. - if len(obj_list) >= 512 and _is_number_sequence(obj_list): - wav_candidates.append(obj_list) - return - - # Otherwise, recurse into elements (but avoid descending into huge numeric lists). - for item in obj_list: - if isinstance(item, list) and len(item) >= 512 and _is_number_sequence(item): # type: ignore[arg-type] - wav_candidates.append(item) - continue - _scan(item, depth + 1) - return - - _scan(ref_audio) - if not sr_candidates: - raise TypeError(f"ref_audio missing sample_rate: {_summarize(ref_audio)}") - sr = int(sr_candidates[0]) - - def _wav_len(x: object) -> int: - try: - if isinstance(x, np.ndarray): - return int(x.size) - if isinstance(x, torch.Tensor): - return int(x.numel()) - if isinstance(x, list): - return int(len(x)) - except Exception: - pass - return 0 - - if not wav_candidates: - raise TypeError(f"ref_audio missing waveform: {_summarize(ref_audio)}") - wav_obj = max(wav_candidates, key=_wav_len) - - def _to_np(x: object) -> np.ndarray: - if isinstance(x, np.ndarray): - return x.astype(np.float32).reshape(-1) - if isinstance(x, torch.Tensor): - return x.detach().to("cpu").float().contiguous().numpy().reshape(-1) - if isinstance(x, dict) and x.get("__ndarray__") and "data" in x and "dtype" in x and "shape" in x: - data = x["data"] - dtype = x["dtype"] - shape = x["shape"] - if isinstance(data, (bytes, bytearray, memoryview)): - return np.frombuffer(data, dtype=dtype).reshape(shape).astype(np.float32).reshape(-1) - if isinstance(x, list): - # list of numbers - if len(x) >= 2 and _is_number_sequence(x): # type: ignore[arg-type] - return np.asarray(x, dtype=np.float32).reshape(-1) - # list of chunks - parts: list[np.ndarray] = [] - for part in x: - if isinstance(part, (np.ndarray, torch.Tensor, list)): - parts.append(_to_np(part)) - if parts: - return np.concatenate(parts, axis=0) - raise TypeError(f"Unsupported waveform type: {type(x)}") - - wav_np = _to_np(wav_obj) - if wav_np.size < 1024: - raise ValueError(f"ref_audio waveform too short: {wav_np.size} samples") - return wav_np, sr - - def _extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor: - # vLLM workers do not automatically move arbitrary torch.nn.Modules to - # CUDA. Ensure the speaker encoder is on the same device/dtype as the - # main model before running it. - dev = next(self.parameters()).device - try: - spk_param = next(self.speaker_encoder.parameters()) - if spk_param.device != dev or spk_param.dtype != torch.bfloat16: - self.speaker_encoder.to(device=dev, dtype=torch.bfloat16) - except StopIteration: - pass - # Resample to 24kHz for speaker encoder. - target_sr = int(getattr(self.config.speaker_encoder_config, "sample_rate", 24000)) - if sr != target_sr: - resampler = self._get_resampler(int(sr), target_sr) - wav = resampler.resample(wav.astype(np.float32), orig_sr=int(sr)) - sr = target_sr - - # Follow official implementation: mel_spectrogram expects 24kHz. Move - # the waveform first so STFT/mel computation stays on the model device - # instead of materializing a CPU mel tensor and copying it per request. - wav_tensor = torch.from_numpy(wav).to(device=dev, dtype=torch.float32).unsqueeze(0) - mel_basis, hann_window = self._get_speaker_mel_buffers( - device=dev, - sampling_rate=24000, - n_fft=1024, - num_mels=128, - win_size=1024, - fmin=0, - fmax=12000, - ) - mels = mel_spectrogram( - wav_tensor, - n_fft=1024, - num_mels=128, - sampling_rate=24000, - hop_size=256, - win_size=1024, - fmin=0, - fmax=12000, - mel_basis=mel_basis, - hann_window=hann_window, - ).transpose(1, 2) - spk = self.speaker_encoder(mels.to(dtype=torch.bfloat16))[0] - return spk.to(dtype=torch.bfloat16) - - def _get_speaker_mel_buffers( + @torch.inference_mode() + def preprocess_batch( self, *, + req_ids: list[str], + model_intermediate_buffer: dict[str, dict[str, Any]], device: torch.device, - sampling_rate: int, - n_fft: int, - num_mels: int, - win_size: int, - fmin: int, - fmax: int | None, - ) -> tuple[torch.Tensor, torch.Tensor]: - cache = getattr(self, "_speaker_mel_buffer_cache", None) - if cache is None: - cache = {} - self._speaker_mel_buffer_cache = cache - key = (str(device), int(sampling_rate), int(n_fft), int(num_mels), int(win_size), int(fmin), fmax) - cached = cache.get(key) - if cached is not None: - return cached - mel_basis = mel_filter_bank( - sr=int(sampling_rate), - n_fft=int(n_fft), - n_mels=int(num_mels), - fmin=int(fmin), - fmax=fmax, - ).to(device) - hann_window = torch.hann_window(int(win_size), device=device) - cached = (mel_basis, hann_window) - cache[key] = cached - return cached + ) -> None: + """Delegate batched preprocess to :class:`Qwen3TTSPromptEmbedsBuilder`.""" + self._prompt_builder.preprocess_batch( + req_ids=req_ids, + model_intermediate_buffer=model_intermediate_buffer, + device=device, + ) def _encode_ref_audio_batch(self, wavs: list[np.ndarray], sr: int, *, device: torch.device) -> list[torch.Tensor]: fe = self._encoder_feature_extractor @@ -1449,765 +883,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): for code, mask in zip(audio_codes, padding_mask) ] - @staticmethod - def _make_ref_audio_cache_key(wav: np.ndarray, sr: int) -> str: - wav_f32 = wav.astype(np.float32, copy=False).reshape(-1) - h = hashlib.sha1() - h.update(int(sr).to_bytes(4, byteorder="little", signed=False)) - h.update(int(wav_f32.size).to_bytes(8, byteorder="little", signed=False)) - h.update(wav_f32.tobytes(order="C")) - return h.hexdigest() - - def _get_ref_audio_artifacts(self, cache_key: str) -> dict[str, torch.Tensor | bool] | None: - entry = self._ref_audio_artifact_cache.get(cache_key) - if entry is None: - return None - # Touch for LRU. - self._ref_audio_artifact_cache.move_to_end(cache_key, last=True) - return entry - - def _put_ref_audio_artifacts( - self, - cache_key: str, - *, - ref_code: torch.Tensor | None = None, - ref_spk_embedding: torch.Tensor | None = None, - ) -> None: - if not cache_key: - return - entry = self._ref_audio_artifact_cache.get(cache_key) - if entry is None: - entry = {} - if isinstance(ref_code, torch.Tensor): - entry["ref_code"] = ref_code.detach().to("cpu", dtype=torch.long).contiguous() - if isinstance(ref_spk_embedding, torch.Tensor): - entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=torch.bfloat16).reshape(-1) - if not entry: - return - self._ref_audio_artifact_cache[cache_key] = entry - self._ref_audio_artifact_cache.move_to_end(cache_key, last=True) - while len(self._ref_audio_artifact_cache) > self._ref_audio_artifact_cache_max_entries: - self._ref_audio_artifact_cache.popitem(last=False) - - @staticmethod - def _first_value(value: object, default: object = None) -> object: - if isinstance(value, list): - return value[0] if value else default - return value if value is not None else default - - @staticmethod - def _coerce_ref_code_tensor(value: object, *, device: torch.device) -> torch.Tensor | None: - value = Qwen3TTSTalkerForConditionalGeneration._first_value(value) - if value is None: - return None - if isinstance(value, torch.Tensor): - ref_code = value - elif isinstance(value, np.ndarray): - ref_code = torch.from_numpy(value) - elif isinstance(value, list) and value: - ref_code = torch.as_tensor(value, dtype=torch.long) - else: - return None - if ref_code.ndim == 3: - ref_code = ref_code[0] - if ref_code.ndim != 2 or ref_code.numel() == 0: - return None - return ref_code.to(device=device, dtype=torch.long).contiguous() - - @staticmethod - def _coerce_token_ids(value: object, *, device: torch.device) -> torch.Tensor | None: - value = Qwen3TTSTalkerForConditionalGeneration._first_value(value) - if value is None: - return None - if isinstance(value, torch.Tensor): - ids = value - elif isinstance(value, np.ndarray): - ids = torch.from_numpy(value) - elif isinstance(value, list) and value and all(isinstance(v, (int, np.integer)) for v in value): - ids = torch.tensor(value, dtype=torch.long) - else: - return None - if ids.ndim == 1: - ids = ids.unsqueeze(0) - if ids.ndim != 2 or ids.numel() == 0: - return None - return ids.to(device=device, dtype=torch.long).contiguous() - - @staticmethod - def _voice_clone_prompt_dict(raw: object) -> dict[str, object] | None: - raw = Qwen3TTSTalkerForConditionalGeneration._first_value(raw) - if isinstance(raw, dict): - return raw - if isinstance(raw, list) and raw and isinstance(raw[0], dict): - return raw[0] - return None - - @staticmethod - def _has_ref_code_like(value: object) -> bool: - value = Qwen3TTSTalkerForConditionalGeneration._first_value(value) - if isinstance(value, torch.Tensor): - return value.numel() > 0 and value.ndim in (2, 3) - if isinstance(value, np.ndarray): - return value.size > 0 and value.ndim in (2, 3) - if isinstance(value, list): - return bool(value) - return False - - @staticmethod - def _needs_initial_prompt_preprocess(info_dict: dict[str, Any]) -> bool: - embed = info_dict.get("embed") - return not (isinstance(embed, dict) and "prefill" in embed) - - def _needs_batched_ref_code(self, info_dict: dict[str, Any]) -> bool: - if not self._needs_initial_prompt_preprocess(info_dict): - return False - if self._first_value(info_dict.get("task_type"), "CustomVoice") != "Base": - return False - codes = info_dict.get("codes") - if isinstance(codes, dict) and self._has_ref_code_like(codes.get(_PRECOMPUTED_REF_CODE_KEY)): - return False - - voice_clone_prompt = self._voice_clone_prompt_dict(info_dict.get("voice_clone_prompt")) - if voice_clone_prompt is not None: - if self._has_ref_code_like(voice_clone_prompt.get("ref_code")): - return False - icl_flag = self._first_value(voice_clone_prompt.get("icl_mode")) - if isinstance(icl_flag, bool) and not icl_flag: - return False - - xvec_only = self._first_value(info_dict.get("x_vector_only_mode"), False) - if isinstance(xvec_only, bool) and xvec_only: - return False - - ref_audio_list = info_dict.get("ref_audio") - return isinstance(ref_audio_list, list) and bool(ref_audio_list) - - @torch.inference_mode() - def preprocess_batch( - self, - *, - req_ids: list[str], - model_intermediate_buffer: dict[str, dict[str, Any]], - device: torch.device, - ) -> None: - """Batch Base voice-clone ref-audio codec extraction for current prefill requests.""" - pending_text: list[tuple[dict[str, Any], str]] = [] - pending_ref_text: list[tuple[dict[str, Any], str]] = [] - groups: dict[int, list[tuple[dict[str, Any], np.ndarray, int]]] = {} - for req_id in req_ids: - info_dict = model_intermediate_buffer.get(req_id) - if not isinstance(info_dict, dict): - continue - - if ( - self._needs_initial_prompt_preprocess(info_dict) - and self._first_value(info_dict.get("task_type"), "CustomVoice") == "Base" - ): - if _PRECOMPUTED_TEXT_IDS_KEY not in info_dict: - text = self._first_value(info_dict.get("text"), "") - if isinstance(text, str): - pending_text.append((info_dict, self._build_assistant_text(text))) - - if _PRECOMPUTED_REF_IDS_KEY not in info_dict: - ref_text = self._first_value(info_dict.get("ref_text")) - if isinstance(ref_text, str) and ref_text.strip(): - pending_ref_text.append((info_dict, self._build_ref_text(ref_text))) - - if not self._needs_batched_ref_code(info_dict): - continue - ref_audio_list = info_dict.get("ref_audio") - if not isinstance(ref_audio_list, list) or not ref_audio_list: - continue - try: - wav, sr = self._normalize_ref_audio(ref_audio_list[0]) - except Exception: - # Keep the original per-request path responsible for surfacing - # invalid ref_audio errors with its existing messages. - continue - info_dict[_NORMALIZED_REF_AUDIO_KEY] = (wav, sr) - cache_key = self._make_ref_audio_cache_key(wav, sr) - info_dict[_REF_AUDIO_CACHE_KEY] = cache_key - cached = self._get_ref_audio_artifacts(cache_key) - if cached is not None: - cached_ref_code = self._coerce_ref_code_tensor(cached.get("ref_code"), device=device) - if isinstance(cached_ref_code, torch.Tensor): - info_dict.setdefault("codes", {})[_PRECOMPUTED_REF_CODE_KEY] = cached_ref_code - continue - groups.setdefault(int(sr), []).append((info_dict, wav, int(sr))) - - if pending_text or pending_ref_text: - tok_text = self._get_tokenizer() - for items, key in ( - (pending_text, _PRECOMPUTED_TEXT_IDS_KEY), - (pending_ref_text, _PRECOMPUTED_REF_IDS_KEY), - ): - if not items: - continue - try: - tokenized = tok_text([text for _, text in items], padding=False) - except Exception as exc: - logger.debug("Qwen3-TTS batched text tokenization failed; falling back to serial path: %s", exc) - continue - input_ids = tokenized.get("input_ids") if isinstance(tokenized, dict) else None - if not isinstance(input_ids, list) or len(input_ids) != len(items): - continue - for (info_dict, _), ids in zip(items, input_ids, strict=True): - if isinstance(ids, list) and ids: - info_dict[key] = torch.tensor(ids, dtype=torch.long) - - groups = {sr: items for sr, items in groups.items() if len(items) >= 2} - if not groups: - return - - for sr, items in groups.items(): - wavs = [wav for _, wav, _ in items] - try: - ref_codes = self._encode_ref_audio_batch(wavs, int(sr), device=device) - except Exception as exc: - logger.debug("Qwen3-TTS batched ref_code encode failed; falling back to serial path: %s", exc) - continue - for (info_dict, wav, item_sr), ref_code in zip(items, ref_codes, strict=True): - ref_code_t = self._coerce_ref_code_tensor(ref_code, device=device) - if ref_code_t is None: - continue - info_dict.setdefault("codes", {})[_PRECOMPUTED_REF_CODE_KEY] = ref_code_t - info_dict[_NORMALIZED_REF_AUDIO_KEY] = (wav, item_sr) - cache_key = info_dict.get(_REF_AUDIO_CACHE_KEY) - if isinstance(cache_key, str) and cache_key: - self._put_ref_audio_artifacts(cache_key, ref_code=ref_code_t) - - def _generate_icl_prompt( - self, - *, - text_id: torch.Tensor, - ref_id: torch.Tensor, - ref_code: torch.Tensor, - tts_pad_embed: torch.Tensor, - tts_eos_embed: torch.Tensor, - non_streaming_mode: bool, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Ported from official Qwen3TTSForConditionalGeneration.generate_icl_prompt - text_embed = self.text_projection(self.text_embedding(torch.cat([ref_id, text_id], dim=-1))) - text_embed = torch.cat([text_embed, tts_eos_embed], dim=1) - - # codec embed (codec bos + codec) 1 T2 D - codec_embed: list[torch.Tensor] = [] - for i in range(int(self.talker_config.num_code_groups)): - if i == 0: - codec_embed.append(self.embed_input_ids(ref_code[:, :1])) - else: - codec_embed.append(self.code_predictor.get_input_embeddings()[i - 1](ref_code[:, i : i + 1])) - codec_embed_sum = torch.cat(codec_embed, dim=1).sum(1).unsqueeze(0) # [1,T,H] - codec_embed_sum = torch.cat( - [ - self.embed_input_ids( - torch.tensor([[self.talker_config.codec_bos_id]], device=codec_embed_sum.device, dtype=torch.long) - ), - codec_embed_sum, - ], - dim=1, - ) - - text_lens = int(text_embed.shape[1]) - codec_lens = int(codec_embed_sum.shape[1]) - if non_streaming_mode: - # Official non-streaming mode: append the full text conditioning in - # prefill, and use PAD in decode steps. - icl_input_embed = text_embed + self.embed_input_ids( - torch.tensor( - [[self.talker_config.codec_pad_id] * text_lens], - device=codec_embed_sum.device, - dtype=torch.long, - ) - ) - icl_input_embed = torch.cat([icl_input_embed, codec_embed_sum + tts_pad_embed], dim=1) - return icl_input_embed, tts_pad_embed - if text_lens > codec_lens: - return text_embed[:, :codec_lens] + codec_embed_sum, text_embed[:, codec_lens:] - text_embed = torch.cat([text_embed] + [tts_pad_embed] * (codec_lens - text_lens), dim=1) - return text_embed + codec_embed_sum, tts_pad_embed - - def _build_prompt_embeds( - self, - *, - task_type: str, - info_dict: dict[str, Any], - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int | None, torch.Tensor | None]: - text = (info_dict.get("text") or [""])[0] - language = (info_dict.get("language") or ["Auto"])[0] - non_streaming_mode_val = info_dict.get("non_streaming_mode") - if isinstance(non_streaming_mode_val, list): - non_streaming_mode_raw = non_streaming_mode_val[0] if non_streaming_mode_val else None - else: - non_streaming_mode_raw = non_streaming_mode_val - if isinstance(non_streaming_mode_raw, bool): - non_streaming_mode = non_streaming_mode_raw - else: - # Match official inference defaults: - # - CustomVoice/VoiceDesign: non_streaming_mode=True - # - Base: non_streaming_mode=False - non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign") - - # Text ids for assistant template (always). - tok = self._get_tokenizer() - dev = next(self.parameters()).device - input_ids = self._coerce_token_ids(info_dict.pop(_PRECOMPUTED_TEXT_IDS_KEY, None), device=dev) - if input_ids is None: - input_ids = tok(self._build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"].to( - device=dev - ) - - # Optional instruct prefix. - instruct = (info_dict.get("instruct") or [""])[0] - instruct_embed = None - if isinstance(instruct, str) and instruct.strip(): - instruct_ids = tok(self._build_instruct_text(instruct), return_tensors="pt", padding=False)["input_ids"].to( - device=input_ids.device - ) - instruct_embed = self.text_projection(self.text_embedding(instruct_ids)) - - # tts special token embeds (projected into talker hidden). - tts_tokens = torch.tensor( - [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]], - device=input_ids.device, - dtype=input_ids.dtype, - ) - tts_bos_embed, tts_eos_embed, tts_pad_embed = self.text_projection(self.text_embedding(tts_tokens)).chunk( - 3, dim=1 - ) - - # Codec prefill tags. - language_id = None - if isinstance(language, str) and language.lower() != "auto": - language_id = self.talker_config.codec_language_id.get(language.lower()) - # Match official dialect override: - # If language is Chinese/Auto and the selected speaker is a dialect voice, - # set language_id to that dialect to improve code generation stability. - if language_id is None and isinstance(language, str) and language.lower() in ("chinese", "auto"): - speaker_for_dialect = None - if task_type == "CustomVoice": - speaker_for_dialect = (info_dict.get("speaker") or [""])[0] - if isinstance(speaker_for_dialect, str) and speaker_for_dialect.strip(): - spk_is_dialect = getattr(self.talker_config, "spk_is_dialect", None) or {} - dialect = spk_is_dialect.get(speaker_for_dialect.lower()) - if isinstance(dialect, str) and dialect: - language_id = self.talker_config.codec_language_id.get(dialect) - if language_id is None: - codec_prefill_list = [ - [ - self.talker_config.codec_nothink_id, - self.talker_config.codec_think_bos_id, - self.talker_config.codec_think_eos_id, - ] - ] - else: - codec_prefill_list = [ - [ - self.talker_config.codec_think_id, - self.talker_config.codec_think_bos_id, - int(language_id), - self.talker_config.codec_think_eos_id, - ] - ] - - codec_input_0 = self.embed_input_ids( - torch.tensor(codec_prefill_list, device=input_ids.device, dtype=torch.long) - ) - codec_input_1 = self.embed_input_ids( - torch.tensor([[self.talker_config.codec_pad_id, self.talker_config.codec_bos_id]], device=input_ids.device) - ) - - # Speaker embedding/token (task-dependent) - speaker_embed = None - ref_code_len: int | None = None - ref_code_prompt: torch.Tensor | None = None - - def _as_singleton(x: object) -> object: - if isinstance(x, list): - return x[0] if x else None - return x - - def _to_long_tensor(x: object, *, device: torch.device) -> torch.Tensor | None: - x = _as_singleton(x) - if x is None: - return None - if isinstance(x, torch.Tensor): - t = x - elif isinstance(x, np.ndarray): - t = torch.from_numpy(x) - elif isinstance(x, list) and x and all(isinstance(v, (int, np.integer)) for v in x): - t = torch.tensor(x, dtype=torch.long) - else: - return None - if t.ndim == 1: - t = t.unsqueeze(0) - return t.to(device=device, dtype=torch.long) - - def _normalize_voice_clone_prompt(raw: object) -> dict[str, object] | None: - raw = _as_singleton(raw) - if raw is None: - return None - if isinstance(raw, dict): - return raw - # Some callers may pass list[dict] directly. - if isinstance(raw, list) and raw and isinstance(raw[0], dict): - return raw[0] - return None - - if task_type == "Base": - # Base supports voice clone prompt with in-context mode. - xvec_only = bool((info_dict.get("x_vector_only_mode") or [False])[0]) - in_context_mode = not xvec_only - voice_clone_prompt = _normalize_voice_clone_prompt(info_dict.get("voice_clone_prompt")) - ref_audio_wav: np.ndarray | None = None - ref_audio_sr: int | None = None - ref_audio_cache_key: str | None = None - - def _get_ref_audio() -> tuple[np.ndarray, int]: - nonlocal ref_audio_wav, ref_audio_sr - if ref_audio_wav is None or ref_audio_sr is None: - normalized_ref_audio = info_dict.pop(_NORMALIZED_REF_AUDIO_KEY, None) - if ( - isinstance(normalized_ref_audio, tuple) - and len(normalized_ref_audio) == 2 - and isinstance(normalized_ref_audio[0], np.ndarray) - ): - ref_audio_wav = normalized_ref_audio[0] - ref_audio_sr = int(normalized_ref_audio[1]) - return ref_audio_wav, ref_audio_sr - ref_audio_list = info_dict.get("ref_audio") - if not isinstance(ref_audio_list, list) or not ref_audio_list: - raise ValueError("Base requires `ref_audio`.") - ref_audio_wav, ref_audio_sr = self._normalize_ref_audio(ref_audio_list[0]) - return ref_audio_wav, ref_audio_sr - - def _get_ref_audio_cache_key() -> str: - nonlocal ref_audio_cache_key - if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key: - return ref_audio_cache_key - cached_key = info_dict.get(_REF_AUDIO_CACHE_KEY) - if isinstance(cached_key, str) and cached_key: - ref_audio_cache_key = cached_key - return ref_audio_cache_key - wav_np, sr = _get_ref_audio() - ref_audio_cache_key = self._make_ref_audio_cache_key(wav_np, sr) - info_dict[_REF_AUDIO_CACHE_KEY] = ref_audio_cache_key - return ref_audio_cache_key - - def _try_get_ref_audio_cache_key() -> str | None: - nonlocal ref_audio_cache_key - if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key: - return ref_audio_cache_key - cached_key = info_dict.get(_REF_AUDIO_CACHE_KEY) - if isinstance(cached_key, str) and cached_key: - ref_audio_cache_key = cached_key - return ref_audio_cache_key - ref_audio_list = info_dict.get("ref_audio") - if isinstance(ref_audio_list, list) and ref_audio_list: - return _get_ref_audio_cache_key() - return None - - # Speaker cache: only for uploaded (named) speakers - _speaker_cache_key = None - if voice_clone_prompt is None: - _speaker_list = info_dict.get("speaker") - if isinstance(_speaker_list, list) and _speaker_list: - _voice_name = str(_speaker_list[0]).lower() - # Per-mode namespace — xvec and icl produce different artifacts - # for the same voice, so they must not share a cache slot. - _mode = "xvec" if xvec_only else "icl" - _voice_created_at = int((info_dict.get("voice_created_at") or [0])[0]) - _speaker_cache_key = self._speaker_cache.make_cache_key( - _voice_name, - model_type=f"qwen3_tts_{_mode}", - created_at=_voice_created_at, - ) - _cached = self._speaker_cache.get(_speaker_cache_key) - if _cached is not None: - # Transfer cached tensors to current device - ref_code_cached = _cached.get("ref_code") - ref_spk_embed_cached = _cached.get("ref_spk_embedding") - if isinstance(ref_code_cached, torch.Tensor): - ref_code_cached = ref_code_cached.to(device=input_ids.device) - if isinstance(ref_spk_embed_cached, torch.Tensor): - ref_spk_embed_cached = ref_spk_embed_cached.to(device=input_ids.device) - voice_clone_prompt = { - "ref_code": ref_code_cached, - "ref_spk_embedding": ref_spk_embed_cached, - "icl_mode": _cached.get("icl_mode"), - } - _speaker_cache_key = None # hit → don't store again - - # Official implementation may pass `voice_clone_prompt.icl_mode`. - if voice_clone_prompt is not None and "icl_mode" in voice_clone_prompt: - icl_flag = _as_singleton(voice_clone_prompt.get("icl_mode")) - if isinstance(icl_flag, bool): - in_context_mode = icl_flag - xvec_only = not in_context_mode - ref_code = None - if voice_clone_prompt is not None: - ref_code = _as_singleton(voice_clone_prompt.get("ref_code")) - ref_code_t = None - if isinstance(ref_code, torch.Tensor): - ref_code_t = ref_code - elif isinstance(ref_code, np.ndarray): - ref_code_t = torch.from_numpy(ref_code) - if isinstance(ref_code_t, torch.Tensor): - if ref_code_t.ndim == 3: - ref_code_t = ref_code_t[0] - ref_code_t = ref_code_t.to(device=input_ids.device, dtype=torch.long) - ref_code_len = int(ref_code_t.shape[0]) - else: - codes = info_dict.get("codes") - precomputed_ref_code = codes.get(_PRECOMPUTED_REF_CODE_KEY) if isinstance(codes, dict) else None - ref_code_t = self._coerce_ref_code_tensor(precomputed_ref_code, device=input_ids.device) - if isinstance(ref_code_t, torch.Tensor): - ref_code_len = int(ref_code_t.shape[0]) - if ref_code_t is None and in_context_mode: - cache_key = _get_ref_audio_cache_key() - cached = self._get_ref_audio_artifacts(cache_key) - cached_ref_code = None - if isinstance(cached, dict): - cached_ref_code = cached.get("ref_code") - ref_code_t = self._coerce_ref_code_tensor(cached_ref_code, device=input_ids.device) - if isinstance(ref_code_t, torch.Tensor): - ref_code_len = int(ref_code_t.shape[0]) - else: - wav_np, sr = _get_ref_audio() - ref_code_t = self._encode_ref_audio_batch([wav_np], int(sr), device=input_ids.device)[0] - ref_code_len = int(ref_code_t.shape[0]) - self._put_ref_audio_artifacts(cache_key, ref_code=ref_code_t) - if isinstance(ref_code_t, torch.Tensor): - ref_code_prompt = ref_code_t - cache_key = _try_get_ref_audio_cache_key() - if isinstance(cache_key, str) and cache_key: - self._put_ref_audio_artifacts(cache_key, ref_code=ref_code_prompt) - - # Speaker embedding: use prompt embed if provided; otherwise extract from audio. - # NOTE: Do NOT use _as_singleton here — the embedding may be a plain - # float list (from API via msgspec IPC) that _as_singleton would - # destructively unwrap to a single scalar. - spk = None - if voice_clone_prompt is not None: - spk = voice_clone_prompt.get("ref_spk_embedding") - if isinstance(spk, torch.Tensor): - speaker_embed = spk.to(device=input_ids.device, dtype=torch.bfloat16).view(1, 1, -1) - elif isinstance(spk, (list, np.ndarray)): - # Plain list/array from API (survived msgspec IPC serialization). - speaker_embed = torch.tensor(spk, dtype=torch.bfloat16, device=input_ids.device).view(1, 1, -1) - else: - cache_key = _get_ref_audio_cache_key() - cached = self._get_ref_audio_artifacts(cache_key) - cached_spk = None - if isinstance(cached, dict): - cached_spk = cached.get("ref_spk_embedding") - if isinstance(cached_spk, torch.Tensor): - speaker_embed = cached_spk.to(device=input_ids.device, dtype=torch.bfloat16).view(1, 1, -1) - else: - wav_np, sr = _get_ref_audio() - speaker_embed = self._extract_speaker_embedding(wav_np, sr).view(1, 1, -1) - self._put_ref_audio_artifacts(cache_key, ref_spk_embedding=speaker_embed) - - # Cache miss: store extraction result - if _speaker_cache_key is not None and speaker_embed is not None: - self._speaker_cache.put( - _speaker_cache_key, - { - "ref_code": ref_code_prompt.detach().cpu() - if isinstance(ref_code_prompt, torch.Tensor) - else None, - "ref_spk_embedding": speaker_embed.detach().cpu().reshape(-1), - "icl_mode": in_context_mode, - }, - ) - - codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1) - - # Role header (<|im_start|>assistant\n) -> projected text embeds. - role_embed = self.text_projection(self.text_embedding(input_ids[:, :3])) - - codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1) - codec_prefix = codec_prefix + codec_input[:, :-1] - talker_prompt = torch.cat((role_embed, codec_prefix), dim=1) - - if in_context_mode: - # Prefer explicit tokenized `ref_ids` if provided (matches official signature). - ref_ids = self._coerce_token_ids(info_dict.pop(_PRECOMPUTED_REF_IDS_KEY, None), device=input_ids.device) - if ref_ids is None: - ref_ids = _to_long_tensor(info_dict.get("ref_ids"), device=input_ids.device) - if ref_ids is None and voice_clone_prompt is not None: - ref_ids = _to_long_tensor( - voice_clone_prompt.get("ref_ids") or voice_clone_prompt.get("ref_id"), device=input_ids.device - ) - if ref_ids is None: - ref_text = _as_singleton(info_dict.get("ref_text")) - if isinstance(ref_text, str) and ref_text.strip(): - ref_ids = tok( - self._build_ref_text(ref_text), - return_tensors="pt", - padding=False, - )["input_ids"].to(device=input_ids.device) - else: - logger.warning("Base ICL: ref_text/ref_ids missing, falling back to x-vector-only mode.") - in_context_mode = False - if in_context_mode: - icl_input_embed, trailing_text_hidden = self._generate_icl_prompt( - text_id=input_ids[:, 3:-5], - ref_id=ref_ids[:, 3:-2], - ref_code=ref_code_t, # type: ignore[arg-type] - tts_pad_embed=tts_pad_embed, - tts_eos_embed=tts_eos_embed, - non_streaming_mode=non_streaming_mode, - ) - talker_prompt = torch.cat([talker_prompt, icl_input_embed], dim=1) - else: - # First text token (+ codec_bos). - if non_streaming_mode: - # Official non-streaming mode: put the full text into the - # prefill prompt and use PAD for decode steps. - text_all = self.text_projection(self.text_embedding(input_ids[:, 3:-5])) - text_all = torch.cat([text_all, tts_eos_embed], dim=1) - pad_ids = torch.full( - (1, int(text_all.shape[1])), - int(self.talker_config.codec_pad_id), - device=input_ids.device, - dtype=torch.long, - ) - talker_prompt = torch.cat( - [ - talker_prompt, - text_all + self.embed_input_ids(pad_ids), - tts_pad_embed - + self.embed_input_ids( - torch.tensor([[self.talker_config.codec_bos_id]], device=input_ids.device) - ), - ], - dim=1, - ) - trailing_text_hidden = tts_pad_embed - else: - first_text = self.text_projection(self.text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:] - talker_prompt = torch.cat([talker_prompt, first_text], dim=1) - trailing_text_hidden = torch.cat( - ( - self.text_projection(self.text_embedding(input_ids[:, 4:-5])), - tts_eos_embed, - ), - dim=1, - ) - - elif task_type == "CustomVoice": - _speaker_raw = info_dict.get("speaker") or [""] - speaker = ( - ((_speaker_raw[0] if isinstance(_speaker_raw, (list, tuple)) else _speaker_raw) or "").lower().strip() - ) - if not speaker: - raise ValueError("CustomVoice requires additional_information.speaker.") - spk_id_map = {k.lower(): v for k, v in (getattr(self.talker_config, "spk_id", None) or {}).items()} - if speaker not in spk_id_map: - raise ValueError(f"Unsupported speaker: {speaker}") - spk_id = spk_id_map[speaker] - # Keep it at least 1D; embedding on a 0-d tensor can return 1D. - spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long) - spk_embed = self.embed_input_ids(spk_tensor) - if spk_embed.ndim == 1: - spk_embed = spk_embed.view(1, 1, -1) - elif spk_embed.ndim == 2: - spk_embed = spk_embed.view(1, 1, -1) - speaker_embed = spk_embed - codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1) - - role_embed = self.text_projection(self.text_embedding(input_ids[:, :3])) - codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1) - codec_prefix = codec_prefix + codec_input[:, :-1] - talker_prompt = torch.cat((role_embed, codec_prefix), dim=1) - - if non_streaming_mode: - text_all = self.text_projection(self.text_embedding(input_ids[:, 3:-5])) - text_all = torch.cat([text_all, tts_eos_embed], dim=1) - pad_ids = torch.full( - (1, int(text_all.shape[1])), - int(self.talker_config.codec_pad_id), - device=input_ids.device, - dtype=torch.long, - ) - talker_prompt = torch.cat( - [ - talker_prompt, - text_all + self.embed_input_ids(pad_ids), - tts_pad_embed - + self.embed_input_ids( - torch.tensor([[self.talker_config.codec_bos_id]], device=input_ids.device) - ), - ], - dim=1, - ) - trailing_text_hidden = tts_pad_embed - else: - first_text = self.text_projection(self.text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:] - talker_prompt = torch.cat([talker_prompt, first_text], dim=1) - trailing_text_hidden = torch.cat( - ( - self.text_projection(self.text_embedding(input_ids[:, 4:-5])), - tts_eos_embed, - ), - dim=1, - ) - - elif task_type == "VoiceDesign": - # No known speaker identity; only codec tags + text. - codec_input = torch.cat([codec_input_0, codec_input_1], dim=1) - - role_embed = self.text_projection(self.text_embedding(input_ids[:, :3])) - codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1) - codec_prefix = codec_prefix + codec_input[:, :-1] - talker_prompt = torch.cat((role_embed, codec_prefix), dim=1) - - if non_streaming_mode: - text_all = self.text_projection(self.text_embedding(input_ids[:, 3:-5])) - text_all = torch.cat([text_all, tts_eos_embed], dim=1) - pad_ids = torch.full( - (1, int(text_all.shape[1])), - int(self.talker_config.codec_pad_id), - device=input_ids.device, - dtype=torch.long, - ) - talker_prompt = torch.cat( - [ - talker_prompt, - text_all + self.embed_input_ids(pad_ids), - tts_pad_embed - + self.embed_input_ids( - torch.tensor([[self.talker_config.codec_bos_id]], device=input_ids.device) - ), - ], - dim=1, - ) - trailing_text_hidden = tts_pad_embed - else: - first_text = self.text_projection(self.text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:] - talker_prompt = torch.cat([talker_prompt, first_text], dim=1) - trailing_text_hidden = torch.cat( - ( - self.text_projection(self.text_embedding(input_ids[:, 4:-5])), - tts_eos_embed, - ), - dim=1, - ) - else: - raise ValueError(f"Unsupported task_type={task_type}") - - if instruct_embed is not None: - talker_prompt = torch.cat([instruct_embed, talker_prompt], dim=1) - - return ( - talker_prompt.squeeze(0), # [prompt_len, H] - trailing_text_hidden.squeeze(0), # [T, H] - tts_pad_embed.squeeze(0), # [1, H] - ref_code_len, - ref_code_prompt.contiguous() if isinstance(ref_code_prompt, torch.Tensor) else None, - ) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Consume talker weights, and conditionally consume speaker encoder # weights only if they are present in the checkpoint. @@ -2269,6 +944,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): device = self.vllm_config.device_config.device self.encoder.to(device=device, dtype=torch.bfloat16) + self._init_runtime_buffers() + logger.info("Loaded %d weights for Qwen3TTSTalkerForConditionalGeneration", len(loaded)) self._build_stacked_codec_embed() return loaded @@ -2282,16 +959,23 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): device=w.device, dtype=w.dtype ) - def _get_resampler(self, orig_sr: int, target_sr: int) -> AudioResampler: - key = (orig_sr, target_sr) - cache = self._resampler_cache - if key in cache: - cache.move_to_end(key) - return cache[key] - cache[key] = AudioResampler(target_sr=target_sr) - if len(cache) > self._resampler_cache_max: - cache.popitem(last=False) - return cache[key] + @torch.no_grad() + def _init_runtime_buffers(self) -> None: + """Populate request-independent runtime buffers from frozen weights. + + Currently this only computes :attr:`_tts_pad_embed` + (``text_proj(text_emb(tts_pad_token_id))``), which the prefill + prompt builder and every decode step mix into the input embedding + in place of an actual text token. The value depends only on the + ``text_embedding`` / ``text_projection`` weights and is therefore + the same for every request — computing it once here avoids + recomputing (and round-tripping through ``info_dict``) on each + forward call. + """ + device = next(self.parameters()).device + pad_ids = torch.tensor([[int(self.config.tts_pad_token_id)]], device=device, dtype=torch.long) + pad_proj = self.text_projection(self.text_embedding(pad_ids)).reshape(1, -1) + self._tts_pad_embed.copy_(pad_proj.to(device=self._tts_pad_embed.device, dtype=self._tts_pad_embed.dtype)) # -------------------- GPU-side MTP fast-path --------------------