[Refactor][Qwen3-TTS] Extract reusable prompt-embeds builder and make tts_pad_embed a persistent buffer (#3992)

Signed-off-by: Viacheslav Klimkov <vklimkov@nvidia.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

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"],

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
 

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),
     )

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"],

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),
         },
     )
 

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,

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"],