feat: revised the continuous streaming to work

- After several testings, the best way is to
append another prefill_embed to the last KV-cache,
while still continuing the length position index.
- This needs to be reset if the start of the next
round of text is equal or greater than position
50 (4 s.). It is the optimum length to not make
the model confused as we append multiple prefill
sequences in between the chunks.
- The next thing is the speaker embedding. Here,
it is revised to use the waveform sequence of the
previous round to make the prosody consistent. As
from our observation, the speaker embedding affects
not only the identity, but also other acoustics,
as well as prosody.
- With factorization of speaker embedding, the previous
point will be better in the performance. However,
for now, this method will work to allow fully
streamable text-to-speech, where you have multiple
rounds of texts.

src/inference/qwen3_tts_inferencer_onnx.py

@@ -117,6 +117,10 @@ _LOCAL_MAX_SEQ_LEN = 16
 _NUM_CODE_GROUPS = 16
 
 
+# For continuous streaming with multi turn texts
+_LENGTH_RESET_LIMIT_FOR_START_OF_MULTI_TURN = 50
+
+
 # ── CPU sampling helpers ──────────────────────────────────────────────────────
 
 
@@ -208,7 +212,7 @@ def _make_session(
     opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
     opts.enable_cpu_mem_arena = True
     opts.enable_mem_pattern = True
-    opts.log_severity_level = 1
+    opts.log_severity_level = 2
 
     if use_cuda:
         cuda_opts = {
@@ -814,6 +818,7 @@ class Qwen3TTSInferencerONNX:
         self._last_hidden_states_np: Optional[np.ndarray] = None
         self._step_idx = 0
         self._talker_seq_len = 0
+        self._generated_wav = None
 
         # ── Streaming text state ──────────────────────────────────────────────
         self._text_cache = ""
@@ -821,6 +826,7 @@ class Qwen3TTSInferencerONNX:
         self._prefilled = False
         self._text_ended = False
         self._turn_idx = 0
+        self._first_step_next_round = False
 
         # ── Audio buffer ──────────────────────────────────────────────────────
         self._prev_tail: Optional[np.ndarray] = None
@@ -1314,7 +1320,10 @@ class Qwen3TTSInferencerONNX:
     def _build_assistant_text(self) -> str:
         return "<|im_start|>assistant\n"
 
-    def _prefill_embeds(self, audio_ref_path: str, language: str) -> np.ndarray:
+    def _build_ending_assistant_text(self) -> str:
+        return "<|im_end|>\n"
+
+    def _prefill_embeds(self, audio_info: AudioLike, language: str) -> np.ndarray:
         """Build the 9-token prefill embedding for the Talker.
 
         Layout (matches modeling_qwen3_tts.py generate_icl_prompt):
@@ -1330,7 +1339,7 @@ class Qwen3TTSInferencerONNX:
                 else:
                     language_id = self._codec_language_id[language.lower()]
         log.info(f"_prefill_embeds language_id {language_id}")
-        speaker_embed = self.create_voice_clone_spkemb(audio_ref_path)  # [B, 1, 512]
+        speaker_embed = self.create_voice_clone_spkemb(audio_info)  # [B, 1, 512]
         log.info(f"_prefill_embeds speaker_embed {speaker_embed} {speaker_embed.shape}")
 
         codec_prefill_list = np.array(
@@ -1401,18 +1410,115 @@ class Qwen3TTSInferencerONNX:
 
     def prefill(self) -> None:
         """Run the 9-token prefill pass (IOBinding + CUDA graph) and snapshot KV."""
-        inputs_embeds = self._prefill_embeds(self._audio_ref_path, self._language)
-        cache_position = np.arange(self._step_idx, self._step_idx + _TALKER_PREFILL_LEN, dtype=np.int64)
+        if self._step_idx == 0:
+            audio_info = self._audio_ref_path
+        else:
+            wav_history = self._generated_wav[0, 0]
+            log.info(f"prefill cont wav_history {wav_history} {wav_history.shape} {wav_history.dtype}")
+            audio_info = (
+                wav_history,
+                24000,
+            )
+            self._generated_wav = None
+        # use previously generated audio sequence
+        # from previous list of texts with the window limit 50 as reference
+        inputs_embeds = self._prefill_embeds(audio_info, self._language)
+        if self._step_idx == 0:
+            cache_position = np.arange(_TALKER_PREFILL_LEN, dtype=np.int64)
+        else:
+            cache_position = np.arange(
+                self._talker_seq_len, self._talker_seq_len + _TALKER_PREFILL_LEN, dtype=np.int64
+            )
         causal_rows = self._causal_tril[:, :, cache_position, :]  # [1,1,q_len,MAX]
         attn_mask = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, MAX_SEQ_LEN]
         cos_rope = self._cos_rope[:, :, cache_position].copy()
         sin_rope = self._sin_rope[:, :, cache_position].copy()
         self._run_talker_prefill(inputs_embeds, cache_position, attn_mask, cos_rope, sin_rope)
-        self._talker_seq_len = self._step_idx + _TALKER_PREFILL_LEN
+        if self._step_idx == 0:
+            self._talker_seq_len = _TALKER_PREFILL_LEN
+        else:
+            self._talker_seq_len = self._talker_seq_len + _TALKER_PREFILL_LEN
+            self._step_idx = self._talker_seq_len
         # Snapshot the post-prefill KV state for turn resets
         self._kv_talker_prefill_np = self._snapshot_talker_kv()
         self._prefilled = True
-        log.info(f"Prefill done; talker_seq_len={self._talker_seq_len}")
+        log.info(f"Prefill done; talker_seq_len={self._talker_seq_len}; step_idx={self._step_idx}")
+
+    def prefill_cont(self) -> None:
+        """
+        Run the prefill for continuation pass (IOBinding + CUDA graph) and snapshot KV.
+        [text_pad],[token_ids("<im_end>\n")],[text_bos] -> continue with the first text token
+        [codec_eos],[None],[codec_pad] -> continue with codec_bos and generation
+        """
+
+        # text_ids = np.array([[self._tts_pad_token_id, self._tts_bos_token_id]], dtype=np.int64)
+        # text_ids = np.array([[self._tts_pad_token_id]], dtype=np.int64)
+        # log.info(f"prefill_cont text_ids {text_ids} {text_ids.shape}")
+        # outputs = self._text_embed_proj.run(["text_emb_out"], {"text_ids": text_ids})  # 3
+        # # embeds = outputs[0]
+        # # tts_pad_embed, tts_bos_embed = embeds[:, :1], embeds[:, 1:]  # 2 * [1 1 d]
+        # tts_pad_embed = outputs[0]
+        # # log.info(f"prefill_cont tts_bos_embed {tts_bos_embed} {tts_bos_embed.shape} {tts_bos_embed.dtype}")
+        # log.info(f"prefill_cont tts_pad_embed {tts_pad_embed} {tts_pad_embed.shape} {tts_pad_embed.dtype}")
+
+        # # codec_ids = np.array([[self._codec_eos_token_id, self._codec_pad_id]], dtype=np.int64)
+        # codec_ids = np.array([[self._codec_eos_token_id]], dtype=np.int64)
+        # log.info(f"prefill_cont codec_ids  {codec_ids}  {codec_ids.shape}")
+        # outputs = self._talker_codec_embed.run(["codec_emb"], {"codec_ids": codec_ids})
+        # # embeds = outputs[0]
+        # # codec_eos_embed, codec_pad_embed = embeds[:, :1], embeds[:, 1:]  # 2 * [1 1 d]
+        # codec_eos_embed = outputs[0]
+        # log.info(f"prefill_cont codec_eos_embed {codec_eos_embed} {codec_eos_embed.shape} {codec_eos_embed.dtype}")
+        # # log.info(f"prefill_cont codec_pad_embed {codec_pad_embed} {codec_pad_embed.shape} {codec_pad_embed.dtype}")
+
+        # tts_pad_codec_eos_embed = tts_pad_embed + codec_eos_embed
+        # # tts_bos_codec_pad_embed = tts_bos_embed + codec_pad_embed
+
+        # self._tokenize_texts([self._build_ending_assistant_text(), self._build_assistant_text()]),
+        # self._tokenize_texts([self._build_assistant_text()]),
+        # prefix_tokens = np.expand_dims(
+        #     np.array(
+        #         self._tokenize_texts([self._build_ending_assistant_text()]),
+        #         dtype=np.int64,
+        #     ),
+        #     axis=0,
+        # )
+        # log.info(f"prefill_cont prefix_tokens {prefix_tokens} {prefix_tokens.shape}")
+        # outputs = self._text_embed_proj.run(["text_emb_out"], {"text_ids": prefix_tokens})  # 3
+        # _talker_input_embed_role = outputs[0]
+        # log.info(
+        #     f"prefill_cont _talker_input_embed_role {_talker_input_embed_role} {_talker_input_embed_role.shape} {_talker_input_embed_role.dtype}"
+        # )
+
+        # talker_input_embed = _talker_input_embed_role
+        # talker_input_embed = np.concatenate(
+        #     (_talker_input_embed_role, tts_bos_codec_pad_embed), axis=1
+        # )  # 6
+        # (tts_pad_codec_eos_embed, _talker_input_embed_role, tts_bos_codec_pad_embed), axis=1
+        # talker_input_embed = np.concatenate((tts_pad_codec_eos_embed, tts_bos_codec_pad_embed), axis=1)  # 2
+        # talker_input_embed = np.concatenate((tts_pad_codec_eos_embed, _talker_input_embed_role), axis=1)  # 3
+        # log.info(
+        #     f"prefill_cont talker_input_embed {talker_input_embed} {talker_input_embed.shape} {talker_input_embed.dtype}"
+        # )
+
+        # log.info(f"prefill_cont talker_seq_len {self._talker_seq_len} step_idx {self._step_idx}")
+        # for i in range(talker_input_embed.shape[1]):
+        #     inputs_embeds = talker_input_embed[:, i : i + 1]
+        #     cache_position = np.array([self._talker_seq_len], dtype=np.int64)
+        #     causal_rows = self._causal_tril[:, :, cache_position, :]  # [1,1,q_len,MAX]
+        #     attn_mask = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, MAX_SEQ_LEN]
+        #     cos_rope = self._cos_rope[:, :, cache_position].copy()
+        #     sin_rope = self._sin_rope[:, :, cache_position].copy()
+        #     self._run_talker_step(inputs_embeds, cache_position, attn_mask, cos_rope, sin_rope)
+        #     self._talker_seq_len += 1
+        #     self._step_idx += 1
+
+        self.prefill()
+        # Snapshot the post-prefill KV state for turn resets
+        self._kv_talker_prefill_np = self._snapshot_talker_kv()
+        self._prefilled = True
+        self._first_step_next_round = True
+        log.info(f"Prefill cont done; talker_seq_len={self._talker_seq_len}; step_idx={self._step_idx}")
 
     # ── Talker prefill execution (IOBinding + CUDA graph) ───────────────────
 
@@ -1563,13 +1669,15 @@ class Qwen3TTSInferencerONNX:
             return self._last_audio_tokens
 
         # ── Build inputs_embeds ───────────────────────────────────────────────
-        if self._step_idx > 0:
+        if self._step_idx > 0 and not self._first_step_next_round:
             codec_embeds = self._last_first_token_embed + self._last_local_tokens_embed
         else:
             codec_embeds = self._talker_codec_embed.run(
                 ["codec_emb"],
                 {"codec_ids": np.array([[self._codec_bos_id]], dtype=np.int64)},
             )[0]
+            if self._first_step_next_round:
+                self._first_step_next_round = False
 
         if text_token is not None:
             text_ids = np.array([[text_token]], dtype=np.int64)
@@ -1594,7 +1702,6 @@ class Qwen3TTSInferencerONNX:
         #     f"step-idx-{self._step_idx} self._last_hidden_states_np {self._last_hidden_states_np} "
         #     f"{self._last_hidden_states_np.shape} {self._last_hidden_states_np.dtype}"
         # )
-        self._talker_seq_len += 1
 
         # ── Sample first token (CPU) ──────────────────────────────────────────
         history = self._generated_tokens[:, :, 0]
@@ -1603,6 +1710,9 @@ class Qwen3TTSInferencerONNX:
 
         self._last_first_token = first_token
         self._is_stopping = bool(first_token[0] == self._codec_eos_token_id)
+
+        self._talker_seq_len += 1
+        self._step_idx += 1
         if self.is_finished:
             return None
 
@@ -1616,7 +1726,6 @@ class Qwen3TTSInferencerONNX:
         # ── Run Local Talker (15 steps) ───────────────────────────────────────
         self._generate_local_transformer()
 
-        self._step_idx += 1
         return self._last_audio_tokens
 
     # ── Local Talker inner loop ───────────────────────────────────────────────
@@ -1953,13 +2062,18 @@ class Qwen3TTSInferencerONNX:
 
     def _drain_pending_tokens(self) -> List[NDArrayInt]:
         outputs: List[NDArrayInt] = []
+        log.info(f"drain_pending_tokens is_prefilled {self._prefilled} is_finished {self.is_finished}")
         if not self._prefilled:
-            self.prefill()
-            return outputs
-        log.info(f"self._pending_tokens {self._pending_tokens}")
+            if self._step_idx == 0:
+                self.prefill()
+            else:
+                self.prefill_cont()
+        log.info(f"drain_pending_tokens pending_tokens {self._pending_tokens} {len(self._pending_tokens)}")
         while self._pending_tokens and not self.is_finished:
             token = self._pending_tokens.pop(0)
-            log.info(f"pop token {token} -> pending_tokens {self._pending_tokens}")
+            log.info(
+                f"drain_pending_tokens token {token} -> pending_tokens {self._pending_tokens} {len(self._pending_tokens)}"
+            )
             output = self.step(token)
             if output is not None:
                 outputs.append(output)
@@ -1967,18 +2081,24 @@ class Qwen3TTSInferencerONNX:
 
     def push_text(self, text_fragment: str) -> List[NDArrayInt]:
         self._text_cache += text_fragment
-        log.info(f"push text_fragment {text_fragment} -> text_cache {self._text_cache}")
+        log.info(f"push_text text_fragment {text_fragment} -> text_cache {self._text_cache}")
         for segment in self._extract_text_segments(force=False):
-            self._pending_tokens.extend(self._tokenize_texts([segment]))
-            log.info(f"push segment {segment} -> pending_tokens {self._pending_tokens}")
+            tokenized_segment = self._tokenize_texts([segment])
+            log.info(
+                f"push_text segment {segment} {len(segment)} -> tokenized_segment {tokenized_segment} {len(tokenized_segment)}"
+            )
+            self._pending_tokens.extend(tokenized_segment)
+            log.info(f"push_text pending_tokens {self._pending_tokens}")
         return self._drain_pending_tokens()
 
     def end_text(self) -> List[NDArrayInt]:
         self._text_ended = True
+        log.info(f"end_text text_cache {self._text_cache} {len(self._text_cache)}")
         if self._text_cache:
             self._pending_tokens.extend(self._tokenize_texts([self._text_cache]))
             self._text_cache = ""
         self._pending_tokens.extend([np.array([self._tts_eos_token_id], dtype=np.int64)][0])
+        log.info(f"end_text pending_tokens {self._pending_tokens} {len(self._pending_tokens)}")
         return self._drain_pending_tokens()
 
     def drain(self, max_steps: Optional[int] = None) -> List[NDArrayInt]:
@@ -2077,6 +2197,10 @@ class Qwen3TTSInferencerONNX:
         )
 
         self._codec_step_idx += self.chunk_frames
+        if self._generated_wav is not None:
+            self._generated_wav = np.concatenate((self._generated_wav, wav.copy()), axis=-1)
+        else:
+            self._generated_wav = wav.copy()
         return wav
 
     def _overlap_samples(self, wav: NDArrayFloat) -> int:
@@ -2141,7 +2265,7 @@ class Qwen3TTSInferencerONNX:
         self._last_hidden_states_np = None
         self._step_idx = 0
 
-    def reset_turn(self, reset_cache: bool = False) -> None:
+    def reset_turn(self, reset_cache: Optional[bool] = None) -> None:
         """Reset for a new turn.  If ``reset_cache=True``, also clears the prefill KV."""
         self._turn_idx += 1
         self._text_cache = ""
@@ -2149,6 +2273,12 @@ class Qwen3TTSInferencerONNX:
         self._prefilled = False
         self._text_ended = False
         self._prev_tail = None
+        self._is_stopping = False
         self._buffer = []
         self._buffer_len = 0
-        self.reset_generation_state(keep_prefill_cache=not reset_cache)
+        if self._step_idx >= _LENGTH_RESET_LIMIT_FOR_START_OF_MULTI_TURN:
+            log.info(f"reset_turn reset_generation_state step-idx-{self._step_idx}")
+            self._first_step_next_round = False
+            self.reset_generation_state(False)
+        elif reset_cache is not None:
+            self.reset_generation_state(keep_prefill_cache=not reset_cache)

test_qwen3-tts-streaming_onnx.py

@@ -102,23 +102,33 @@ _CODEC_CONFIG_PATH = "./configs/speech_tokenizer_config.json"
 # _AUDIO_REF_PATH = "./audio_ref/female_shadowheart.flac"
 _AUDIO_REF_PATH = "./audio_ref/male_stewie.mp3"
 _OUTPUT_WAV_DIR = "./audio_synth/"
-_LANGUAGE = "russian"
 # _LANGUAGE = "english"
+# _TEXT = [
+#     "Depending on the time,",
+#     "not only accuracy,",
+#     "but also low latency is important.",
+#     "If it is not instant,",
+#     "then, the human interaction is lost",
+#     "We are finally reaching a moment",
+#     "where the technology is fast enough",
+#     "for people to simply communicate.",
+#     "And that is a huge shift",
+#     "for global business",
+# ]
+_LANGUAGE = "russian"
 # _TEXT = "в зависимости от времени не только точность, но и низкая задержка."
 _TEXT = [
-"Проверьте, проверьте, тест, один,",
-"два Это система непрерывного перевода",
-"речи в зависимости от времени не только",
-"но не только точность,",
-"но и но также важна задержка если это не мгновенн>о",
-"тогда то человеческое взаимодействие",
-"теряется мы наконец-то достигаем момента,",
-"когда техн>ология станет достаточно быстрой для того,"
-"чтобы люди могли просто общаться и это огромный с>",
-"двиг для глобального бизнеса"
+    "В зависимости от ситуации,",
+    "важна не только точность,",
+    "но и низкая задержка.",
+    "Если это происходит не мгновенно,",
+    "то теряется эффект живого общения.",
+    "Мы наконец-то достигли момента,",
+    "когда технологии стали достаточно быстрыми,",
+    "чтобы люди могли просто общаться.",
+    "И это огромный сдвиг",
+    "для мирового бизнеса.",
 ]
-# _TEXT = ["This is a test, ", "how about this one?", "Is this also?"]
-# _TEXT = ["This is not a test.", "How about this one?", "Is this also not?"]
 _OUTPUT_SAMPLE_RATE = 24000
 _CHUNK_FRAMES = 4
 _TEMPERATURE = 0.85
@@ -204,7 +214,7 @@ def run_streaming_tts(
     t_start = time.perf_counter()
 
     for i, text in enumerate(text_list):
-        inferencer.reset_turn(reset_cache=False)
+        inferencer.reset_turn(reset_cache=None)
         log.info(f"Synthesising: {text!r}")
 
         # ── Stream text deltas ────────────────────────────────────────────────────