feat: improve the last chunk with dynamic codec

- Previously, the chunk_frames length for the
last chunk of each rounds follow the fixed 4 (320 ms.).
This caused requirement to flush the codec cache
after every round of text due to zero pad right,
which in turns caused annoying pulses between
round of text.
- Now, it has been fixed with a separate onnx session
of codec decoder with dynamic axis on the chunk frames
length, which is without CUDA graph. This allows the
cache to be continuous over all rounds, as all of
the caches of the codec decoder are ring buffer type.
- Now, the generated speech sounds smooth with real-time
streaming processing within each round of text and
across rounds of texts.
- Final polishing will be to slightly put cross-fade
just in case the amplitude jumps too sharp between
rounds.

.gitattributes

@@ -63,3 +63,4 @@ audio_synth/output_1776738826.4899054.wav filter=lfs diff=lfs merge=lfs -text
 audio_synth/output_1776994773.978851.wav filter=lfs diff=lfs merge=lfs -text
 audio_synth/output_1777092578.4959543.wav filter=lfs diff=lfs merge=lfs -text
 audio_synth/output_1777159335.3181114.wav filter=lfs diff=lfs merge=lfs -text
+qwen3-tts_onnx/codec_decoder_model_dynamic_chunks.onnx filter=lfs diff=lfs merge=lfs -text

qwen3-tts_onnx/codec_decoder_model_dynamic_chunks.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c928257877ac48d058bae7a90357f09ca59fa0445822aec70e92f3906fd9c66d
+size 456883759

src/inference/qwen3_tts_inferencer_onnx.py

@@ -366,6 +366,7 @@ class Qwen3TTSInferencerONNX:
     talker_local_lm_head_model_path : str
         Path to ``talker_local_lm_head.onnx``  (batched lm_head for all 15 codebook groups).
     codec_decoder_model_path : str
+    codec_decoder_model_dynamic_chunks_path : str
     speaker_encoder_model_path : str
     talker_codec_embed_model_path : str
     text_embed_proj_model_path : str
@@ -395,6 +396,7 @@ class Qwen3TTSInferencerONNX:
         talker_local_model_step_path: str,
         talker_local_lm_head_model_path: str,
         codec_decoder_model_path: str,
+        codec_decoder_model_dynamic_chunks_path: str,
         speaker_encoder_model_path: str,
         talker_codec_embed_model_path: str,
         text_embed_proj_model_path: str,
@@ -408,11 +410,11 @@ class Qwen3TTSInferencerONNX:
         enable_cuda_graph: bool = True,
         chunk_frames: int = 4,
         num_threads: int = 4,
-        temperature: float = 0.85,
-        top_p: float = 0.8,
+        temperature: float = 0.75,
+        top_p: float = 0.85,
         top_k: int = 50,
-        repetition_penalty: float = 1.9,
-        repetition_window: int = 50,
+        repetition_penalty: float = 9.5,
+        repetition_window: int = 75,
     ) -> None:
 
         self._use_cuda = use_cuda
@@ -449,9 +451,13 @@ class Qwen3TTSInferencerONNX:
         log.info(f"  local lm_head   <- {talker_local_lm_head_model_path}")
         _local_lm_head_sess = _sess(talker_local_lm_head_model_path, graph=self._enable_cuda_graph)
 
-        # Non-graphed utility sessions
+        # Codec decoder session: IOBinding + CUDA graph (chunk_frames fixed to 4 [320 ms.])
         log.info(f"  codec decoder   <- {codec_decoder_model_path}")
         self._codec_decoder_sess = _sess(codec_decoder_model_path, graph=self._enable_cuda_graph)
+
+        # Non-graphed utility sessions
+        log.info(f"  codec decoder dynamic chunks   <- {codec_decoder_model_dynamic_chunks_path}")
+        self._codec_decoder_dynamic_chunks = _sess(codec_decoder_model_dynamic_chunks_path, graph=False)
         log.info(f"  speaker encoder <- {speaker_encoder_model_path}")
         self._speaker_encoder = _sess(speaker_encoder_model_path, graph=False)
         log.info(f"  codec embed     <- {talker_codec_embed_model_path}")
@@ -625,10 +631,6 @@ class Qwen3TTSInferencerONNX:
         _causal_tril_local = np.tril(np.ones((self._num_code_groups, self._num_code_groups), dtype=bool))
         self._causal_tril_local = _causal_tril_local.reshape(1, 1, self._num_code_groups, self._num_code_groups)
 
-        # Define cache_write mask for each position in the Local prefill and steps, used to update KV-cache
-        # self._local_cache_write_mask_list = self._make_local_cache_write_mask(batch_size=1)  # 15 cache write masks
-        # For each, [B, H, 16, D], with varying 1/0 depending on the position
-
         # Define cos and sin table for Local RoPE embedding
         self._cos_rope_local, self._sin_rope_local = self._build_rope_tables_numpy(
             rope_dim=self._code_predictor_config.head_dim,
@@ -975,6 +977,15 @@ class Qwen3TTSInferencerONNX:
         for ov in self._kv_local_ov:
             _copy_numpy_to_ortvalue(z, ov)
 
+    def _snapshot_codec_kv(self) -> List[np.ndarray]:
+        """Copy all Codec KV device OrtValues to CPU NumPy arrays."""
+        return [ov.numpy().copy() for ov in self._codec_kv_ovs]
+
+    def _restore_codec_kv(self, snapshot: List[np.ndarray]) -> None:
+        """Copy a CPU snapshot back into the Codec KV device OrtValues."""
+        for ov, arr in zip(self._codec_kv_ovs, snapshot):
+            _copy_numpy_to_ortvalue(arr, ov)
+
     def _zero_codec_kv(self) -> None:
         """Zero out all Codec Decoder device buffers and reset the step counter."""
         zero_wav = np.zeros(
@@ -1865,6 +1876,78 @@ class Qwen3TTSInferencerONNX:
         avg_exec_time = exec_time_tot / len(times["codec_decoder"]) * 1000
         log.info(f"\t{n_iter} iter. run codec_decoder in {exec_time_tot:.2f} s. --> {avg_exec_time:.2f} ms. per run.")
 
+        # Codec decoder with dynamic chunk frames
+        n_chunk_frames = 2
+        log.info(f"Warming-up codec decoder with dynamic chunk frames (n={n_chunk_frames})...")
+        self._zero_codec_kv()
+        chunk_tokens = np.random.randint(0, self._local_vocab_size, size=(1, self._num_code_groups, n_chunk_frames))
+        n_iter = n_iter // n_chunk_frames
+        codec_step_idx = 0
+        for n in range(n_iter):
+            pos = np.arange(codec_step_idx, codec_step_idx + n_chunk_frames, dtype=np.int64)
+            cos_rope = self._cos_rope_codec[:, pos].copy()
+            sin_rope = self._sin_rope_codec[:, pos].copy()
+            if codec_step_idx + n_chunk_frames < self._speech_tokenizer_sliding_window:
+                # attention mask goes from the right to the left
+                # we pad left False if the current + chunks still less than sliding window left
+                pad_len = self._speech_tokenizer_sliding_window - (codec_step_idx + n_chunk_frames)
+                # here, basically, we need to pad left because we start from the right
+                # and there still remaining amounts before it attends full sliding window length
+                # so we must not attend to those previous remaining amounts, which is less than pos. 0
+                causal_rows_right = self._causal_tril_codec[:, :, pos, :-pad_len]  # [1,1,q_len,72-accum_len]
+                causal_rows = np.pad(
+                    causal_rows_right, ((0, 0), (0, 0), (0, 0), (pad_len, 0)), mode="constant", constant_values=False
+                )
+            else:
+                # this case where the current + chunks is exactly sliding window length
+                # or more than sliding window length
+                # simply take the last chunks from the causal tril
+                causal_rows = self._causal_tril_codec[:, :, -n_chunk_frames:, :]  # [1,1,q_len,72]
+            attn_mask = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, 72]
+            # use session with dynamic axes as this is the last chunk of this round
+            # and it is less than the default chunk frames 4
+            # copy cache from ortvalue of the cuda-graph static session
+            hidden_state_cache = self._codec_hidden_cache_ov.numpy()
+            pre_conv_hidden_state_cache = self._codec_pre_conv_cache_ov.numpy()
+            kv_cache = self._snapshot_codec_kv()
+            # set with the input names
+            feed = {
+                "codes": chunk_tokens,
+                "hidden_state_cache": hidden_state_cache,
+                "pre_conv_hidden_state_cache": pre_conv_hidden_state_cache,
+                "attention_mask": attn_mask,
+                "cos_rope": cos_rope,
+                "sin_rope": sin_rope,
+            }
+            for i in range(self._speech_tokenizer_num_hidden_layers):
+                feed[f"past_key_{i}"] = kv_cache[2 * i]
+                feed[f"past_value_{i}"] = kv_cache[2 * i + 1]
+            # set output names
+            output_names = ["wav", "current_hidden_state_cache", "current_pre_conv_hidden_state_cache"]
+            for i in range(self._speech_tokenizer_num_hidden_layers):
+                output_names.extend([f"present_key_{i}", f"present_value_{i}"])
+            # run the session with dynamic axes on the chunk frame
+            t0 = time.perf_counter()
+            outputs = self._codec_decoder_dynamic_chunks.run(output_names, feed)
+            times["codec_decoder_dynamic"].append((time.perf_counter() - t0) * 1000)  # ms
+            # get the output
+            hidden_state_cache, pre_conv_hidden_state_cache, past_key_values = (
+                outputs[1],
+                outputs[2],
+                outputs[3:],
+            )
+            # update the bound ort cache for the cuda-graph static session
+            _copy_numpy_to_ortvalue(hidden_state_cache, self._codec_hidden_cache_ov)
+            _copy_numpy_to_ortvalue(pre_conv_hidden_state_cache, self._codec_pre_conv_cache_ov)
+            self._restore_codec_kv(past_key_values)
+            codec_step_idx += n_chunk_frames
+        exec_time_tot = sum(times["codec_decoder_dynamic"]) / 1000
+        avg_exec_time = exec_time_tot / len(times["codec_decoder_dynamic"]) * 1000
+        log.info(
+            f"\t{n_iter} iter. run codec_decoder_dynamic in {exec_time_tot:.2f} s. "
+            f"--> {avg_exec_time:.2f} ms. per run."
+        )
+
         # Summary log
         log.info("Warmup latency summary:")
         for k, v in times.items():
@@ -2011,31 +2094,24 @@ class Qwen3TTSInferencerONNX:
     def _process_frames_to_audio(self, num_frames: int) -> NDArrayFloat:
         """Run Codec Decoder via IOBinding + CUDA graph."""
         chunk_tokens = self._consume_frames(num_frames)
-        len_pad_right = 0
-        log.info(f"codec-idx-{self._codec_step_idx} chunk_tokens {chunk_tokens} {chunk_tokens.shape}")
         if chunk_tokens.shape[-1] < self.chunk_frames:
-            len_pad_right = self.chunk_frames - chunk_tokens.shape[-1]
-            chunk_tokens = np.concatenate(
-                (
-                    chunk_tokens,
-                    np.zeros((chunk_tokens.shape[0], chunk_tokens.shape[1], len_pad_right)).astype(np.int64),
-                ),
-                axis=-1,
-            )
-            log.info(
-                f"codec-idx-{self._codec_step_idx} len_pad_right {len_pad_right} "
-                f"chunk_tokens {chunk_tokens} {chunk_tokens.shape}"
-            )
+            n_chunk_frames = chunk_tokens.shape[-1]
+        else:
+            n_chunk_frames = self.chunk_frames
+        log.info(
+            f"codec-idx-{self._codec_step_idx} chunk_tokens {chunk_tokens} {chunk_tokens.shape} "
+            f"n_chunk_frames {n_chunk_frames}"
+        )
 
         # cache_position for the decoder transformer
         # It starts from 0 and grows. The model handles the sliding window internally.
-        pos = np.arange(self._codec_step_idx, self._codec_step_idx + self.chunk_frames, dtype=np.int64)
+        pos = np.arange(self._codec_step_idx, self._codec_step_idx + n_chunk_frames, dtype=np.int64)
         cos_rope = self._cos_rope_codec[:, pos].copy()
         sin_rope = self._sin_rope_codec[:, pos].copy()
-        if self._codec_step_idx + self.chunk_frames < self._speech_tokenizer_sliding_window:
+        if self._codec_step_idx + n_chunk_frames < self._speech_tokenizer_sliding_window:
             # attention mask goes from the right to the left
             # we pad left False if the current + chunks still less than sliding window left
-            pad_len = self._speech_tokenizer_sliding_window - (self._codec_step_idx + self.chunk_frames)
+            pad_len = self._speech_tokenizer_sliding_window - (self._codec_step_idx + n_chunk_frames)
             # here, basically, we need to pad left because we start from the right
             # and there still remaining amounts before it attends full sliding window length
             # so we must not attend to those previous remaining amounts, which is less than pos. 0
@@ -2047,24 +2123,57 @@ class Qwen3TTSInferencerONNX:
             # this case where the current + chunks is exactly sliding window length
             # or more than sliding window length
             # simply take the last chunks from the causal tril
-            causal_rows = self._causal_tril_codec[:, :, -self.chunk_frames :, :]  # [1,1,q_len,72]
+            causal_rows = self._causal_tril_codec[:, :, -n_chunk_frames:, :]  # [1,1,q_len,72]
         attn_mask = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, 72]
 
-        # Run the wired graph
-        # chunk_tokens need to be .copy(), otherwise it messed up if we bind input and/or use CUDA graph
-        wav = self._run_codec_decoder(chunk_tokens.copy(), attn_mask, cos_rope, sin_rope).copy()
-        # new memory address, so that output wav not only the last repeated
+        if n_chunk_frames == self.chunk_frames:
+            # Run the wired graph
+            # chunk_tokens need to be .copy(), otherwise it messed up if we bind input and/or use CUDA graph
+            wav = self._run_codec_decoder(chunk_tokens.copy(), attn_mask, cos_rope, sin_rope).copy()
+            # new memory address, so that output wav not only the last repeated
+        else:
+            # use session with dynamic axes as this is the last chunk of this round
+            # and it is less than the default chunk frames 4
+            # copy cache from ortvalue of the cuda-graph static session
+            hidden_state_cache = self._codec_hidden_cache_ov.numpy()
+            pre_conv_hidden_state_cache = self._codec_pre_conv_cache_ov.numpy()
+            kv_cache = self._snapshot_codec_kv()
+            # set with the input names
+            feed = {
+                "codes": chunk_tokens,
+                "hidden_state_cache": hidden_state_cache,
+                "pre_conv_hidden_state_cache": pre_conv_hidden_state_cache,
+                "attention_mask": attn_mask,
+                "cos_rope": cos_rope,
+                "sin_rope": sin_rope,
+            }
+            for i in range(self._speech_tokenizer_num_hidden_layers):
+                feed[f"past_key_{i}"] = kv_cache[2 * i]
+                feed[f"past_value_{i}"] = kv_cache[2 * i + 1]
+            # set output names
+            output_names = ["wav", "current_hidden_state_cache", "current_pre_conv_hidden_state_cache"]
+            for i in range(self._speech_tokenizer_num_hidden_layers):
+                output_names.extend([f"present_key_{i}", f"present_value_{i}"])
+            # run the session with dynamic axes on the chunk frame
+            outputs = self._codec_decoder_dynamic_chunks.run(output_names, feed)
+            # get the output
+            wav, hidden_state_cache, pre_conv_hidden_state_cache, past_key_values = (
+                outputs[0],
+                outputs[1],
+                outputs[2],
+                outputs[3:],
+            )
+            # update the bound ort cache for the cuda-graph static session
+            _copy_numpy_to_ortvalue(hidden_state_cache, self._codec_hidden_cache_ov)
+            _copy_numpy_to_ortvalue(pre_conv_hidden_state_cache, self._codec_pre_conv_cache_ov)
+            self._restore_codec_kv(past_key_values)
 
-        # The decoder output already discards the upsampled left-context frames (25 frames)
-        # Just crop the right-pad region from the waveform if chunk_tokens padded
-        if len_pad_right > 0:
-            wav = wav[..., : -(len_pad_right * self._speech_tokenizer_decoder_total_upsample)]
         log.info(
             f"codec-idx-{self._codec_step_idx} wav {wav} "
             f"{np.min(wav)} {np.mean(wav)} {np.std(wav)} {np.max(wav)} {wav.shape}"
         )
 
-        self._codec_step_idx += self.chunk_frames
+        self._codec_step_idx += n_chunk_frames
         if self._ref_wav_history_list is not None:
             self._ref_wav_history_list.append(wav.copy())
         else:

test_qwen3-tts-streaming_onnx.py

@@ -319,6 +319,7 @@ def main() -> None:
 
     # Other model paths
     codec_decoder_path = _onnx_path(args.onnx_dir, f"codec_decoder_model{suffix}")
+    codec_decoder_dynamic_chunks_path = _onnx_path(args.onnx_dir, f"codec_decoder_model_dynamic_chunks{suffix}")
     speaker_encoder_path = _onnx_path(args.onnx_dir, f"speaker_encoder_model{suffix}")
     talker_codec_embed_path = _onnx_path(args.onnx_dir, f"talker_codec_embed_model{suffix}")
     text_embed_proj_path = _onnx_path(args.onnx_dir, f"text_embed_proj_model{suffix}")
@@ -331,6 +332,7 @@ def main() -> None:
         talker_local_step_path,
         talker_local_lm_head_path,
         codec_decoder_path,
+        codec_decoder_dynamic_chunks_path,
         speaker_encoder_path,
         talker_codec_embed_path,
         text_embed_proj_path,
@@ -351,6 +353,7 @@ def main() -> None:
         talker_local_model_step_path=talker_local_step_path,
         talker_local_lm_head_model_path=talker_local_lm_head_path,
         codec_decoder_model_path=codec_decoder_path,
+        codec_decoder_model_dynamic_chunks_path=codec_decoder_dynamic_chunks_path,
         speaker_encoder_model_path=speaker_encoder_path,
         talker_codec_embed_model_path=talker_codec_embed_path,
         text_embed_proj_model_path=text_embed_proj_path,