Update nano-tts

gpt2_decoder.py

@@ -194,13 +194,24 @@ class NanoGPT2Attention(nn.Module):
         key = key.transpose(1, 2)
         value = value.transpose(1, 2)
         mask = None
+        query_attention_mask = None
         if attention_mask is not None:
+            query_length = query.shape[-2]
+            key_length = key.shape[-2]
             mask = self._causal_attention_mask(
                 attention_mask=attention_mask,
-                query_length=query.shape[-2],
-                key_length=key.shape[-2],
+                query_length=query_length,
+                key_length=key_length,
                 device=query.device,
             )
+            query_attention_mask = attention_mask[:, -query_length:].to(dtype=torch.bool, device=query.device)
+            if not bool(query_attention_mask.all()):
+                # SDPA can produce NaNs when a query row is fully masked. For padded query positions,
+                # keep a single aligned key visible, then zero the query output after attention.
+                mask = mask.expand(query.shape[0], -1, -1, -1).clone()
+                invalid_batch, invalid_query = torch.nonzero(~query_attention_mask, as_tuple=True)
+                aligned_key = invalid_query + max(key_length - query_length, 0)
+                mask[invalid_batch, :, invalid_query, aligned_key] = True
         output = torch.nn.functional.scaled_dot_product_attention(
             query,
             key,
@@ -209,6 +220,8 @@ class NanoGPT2Attention(nn.Module):
             dropout_p=self.attn_dropout if self.training else 0.0,
             is_causal=mask is None,
         )
+        if query_attention_mask is not None and not bool(query_attention_mask.all()):
+            output = output.masked_fill(~query_attention_mask[:, None, :, None], 0.0)
         return output.transpose(1, 2).contiguous()
 
     def _flash_attention(

modeling_nanotts_global_local.py

@@ -4,10 +4,11 @@ from __future__ import annotations
 import json
 import logging
 import os
+import time
 from contextlib import nullcontext
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Iterator, Optional, Sequence, Union
 
 import numpy as np
 import torch
@@ -666,6 +667,49 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         )
         return resolved_batch_size
 
+    @staticmethod
+    def _resolve_requested_batch_size_limit(requested_batch_size: Optional[int]) -> Optional[int]:
+        if requested_batch_size is None:
+            return None
+        resolved_batch_size = int(requested_batch_size)
+        if resolved_batch_size <= 0:
+            return None
+        return max(1, resolved_batch_size)
+
+    def _resolve_effective_voice_clone_batch_sizes(
+        self,
+        *,
+        resolved_device: torch.device,
+        chunk_count: int,
+        max_memory_per_sample_gb: float,
+        requested_tts_max_batch_size: Optional[int] = None,
+        requested_codec_max_batch_size: Optional[int] = None,
+        realtime_streaming: bool = False,
+    ) -> tuple[int, int]:
+        effective_tts_batch_size = self._resolve_voice_clone_chunk_batch_size(
+            resolved_device=resolved_device,
+            chunk_count=chunk_count,
+            max_memory_per_sample_gb=max_memory_per_sample_gb,
+        )
+        requested_tts_limit = self._resolve_requested_batch_size_limit(requested_tts_max_batch_size)
+        requested_codec_limit = self._resolve_requested_batch_size_limit(requested_codec_max_batch_size)
+
+        if requested_tts_limit is not None:
+            effective_tts_batch_size = min(effective_tts_batch_size, requested_tts_limit)
+        if realtime_streaming and requested_codec_limit is not None:
+            effective_tts_batch_size = min(effective_tts_batch_size, requested_codec_limit)
+
+        effective_tts_batch_size = max(1, min(int(chunk_count), int(effective_tts_batch_size)))
+
+        if realtime_streaming:
+            effective_codec_batch_size = effective_tts_batch_size
+        elif requested_codec_limit is None:
+            effective_codec_batch_size = 1
+        else:
+            effective_codec_batch_size = max(1, min(int(requested_codec_limit), int(effective_tts_batch_size)))
+
+        return int(effective_tts_batch_size), int(effective_codec_batch_size)
+
     def _generate_audio_token_ids_with_fallback(
         self,
         *,
@@ -740,6 +784,67 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         )
         return self._extract_waveform_and_sample_rate(decoded, fallback_sample_rate=target_sample_rate)
 
+    def _decode_audio_token_id_batch_to_waveforms(
+        self,
+        *,
+        audio_tokenizer,
+        audio_token_id_batches: Sequence[torch.LongTensor],
+        target_sample_rate: int,
+        effective_nq: int,
+        resolved_device: torch.device,
+    ) -> tuple[list[torch.FloatTensor], int]:
+        if not audio_token_id_batches:
+            return [], target_sample_rate
+
+        if len(audio_token_id_batches) == 1:
+            waveform, sample_rate = self._decode_audio_token_ids_to_waveform(
+                audio_tokenizer=audio_tokenizer,
+                audio_token_ids=audio_token_id_batches[0],
+                target_sample_rate=target_sample_rate,
+                effective_nq=effective_nq,
+                resolved_device=resolved_device,
+            )
+            return [waveform], sample_rate
+
+        decode_codes = [
+            self._prepare_audio_codes_for_decode(audio_token_ids.to(resolved_device), nq=effective_nq)
+            for audio_token_ids in audio_token_id_batches
+        ]
+        batch_decode_fn = getattr(audio_tokenizer, "batch_decode", None)
+        if batch_decode_fn is None:
+            raise AttributeError("audio_tokenizer must provide a batch_decode method.")
+
+        try:
+            with self._audio_tokenizer_inference_context(audio_tokenizer, resolved_device):
+                decode_output = batch_decode_fn(
+                    decode_codes,
+                    num_quantizers=effective_nq,
+                    chunk_duration=None,
+                )
+            return self._extract_batch_waveforms_and_sample_rate(
+                decode_output,
+                fallback_sample_rate=target_sample_rate,
+                batch_size=len(audio_token_id_batches),
+            )
+        except Exception:
+            logging.warning(
+                "batched audio decode failed; falling back to per-chunk decode for batch_size=%d",
+                len(audio_token_id_batches),
+                exc_info=True,
+            )
+            waveform_rows: list[torch.FloatTensor] = []
+            sample_rate = target_sample_rate
+            for audio_token_ids in audio_token_id_batches:
+                waveform_row, sample_rate = self._decode_audio_token_ids_to_waveform(
+                    audio_tokenizer=audio_tokenizer,
+                    audio_token_ids=audio_token_ids,
+                    target_sample_rate=target_sample_rate,
+                    effective_nq=effective_nq,
+                    resolved_device=resolved_device,
+                )
+                waveform_rows.append(waveform_row)
+            return waveform_rows, sample_rate
+
     def _build_generation_row(
         self,
         batch_size: int,
@@ -756,6 +861,38 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         row[:, :, 1:] = audio_token_ids.unsqueeze(1)
         return row
 
+    @staticmethod
+    def _compute_stream_lead_seconds(
+        emitted_samples_total: int,
+        sample_rate: int,
+        first_audio_emitted_at: Optional[float],
+    ) -> float:
+        if first_audio_emitted_at is None or sample_rate <= 0:
+            return 0.0
+        elapsed_seconds = max(0.0, time.monotonic() - first_audio_emitted_at)
+        emitted_seconds = float(emitted_samples_total) / float(sample_rate)
+        return emitted_seconds - elapsed_seconds
+
+    @staticmethod
+    def _resolve_stream_decode_frame_budget(
+        *,
+        emitted_samples_total: int,
+        sample_rate: int,
+        first_audio_emitted_at: Optional[float],
+    ) -> int:
+        lead_seconds = NanoTTSGlobalLocalForCausalLM._compute_stream_lead_seconds(
+            emitted_samples_total=emitted_samples_total,
+            sample_rate=sample_rate,
+            first_audio_emitted_at=first_audio_emitted_at,
+        )
+        if first_audio_emitted_at is None or lead_seconds < 0.20:
+            return 1
+        if lead_seconds < 0.55:
+            return 2
+        if lead_seconds < 1.10:
+            return 4
+        return 8
+
     def _sample_next_token(
         self,
         logits: torch.FloatTensor,
@@ -1300,6 +1437,74 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         with self._audio_tokenizer_inference_context(audio_tokenizer, decode_codes.device):
             return batch_decode_fn([decode_codes], num_quantizers=effective_nq, chunk_duration=None)
 
+    def _extract_batch_waveforms_and_sample_rate(
+        self,
+        decode_output: Any,
+        fallback_sample_rate: int,
+        batch_size: int,
+    ) -> tuple[list[torch.FloatTensor], int]:
+        sample_rate = fallback_sample_rate
+        audio = decode_output
+        audio_lengths = None
+
+        for key in ("sample_rate", "sampling_rate"):
+            value = getattr(decode_output, key, None)
+            if value is not None:
+                sample_rate = int(value)
+                break
+        for key in ("waveform", "audio", "wav", "samples"):
+            value = getattr(decode_output, key, None)
+            if value is not None:
+                audio = value
+                break
+        for key in ("audio_lengths", "waveform_lengths", "lengths"):
+            value = getattr(decode_output, key, None)
+            if value is not None:
+                audio_lengths = value
+                break
+
+        if isinstance(decode_output, dict):
+            for key in ("sample_rate", "sampling_rate"):
+                if key in decode_output:
+                    sample_rate = int(decode_output[key])
+                    break
+            for key in ("waveform", "audio", "wav", "samples"):
+                if key in decode_output:
+                    audio = decode_output[key]
+                    break
+            for key in ("audio_lengths", "waveform_lengths", "lengths"):
+                if key in decode_output:
+                    audio_lengths = decode_output[key]
+                    break
+
+        audio_tensor = torch.as_tensor(audio, dtype=torch.float32)
+        if audio_tensor.ndim == 2:
+            audio_tensor = audio_tensor.unsqueeze(0)
+        if audio_tensor.ndim != 3:
+            raise ValueError(f"Expected batched decoded audio with 3 dims, got {tuple(audio_tensor.shape)}")
+        if audio_tensor.shape[0] != int(batch_size):
+            raise ValueError(
+                f"Expected decoded batch size {batch_size}, got audio tensor shape {tuple(audio_tensor.shape)}"
+            )
+
+        if audio_lengths is None:
+            lengths_tensor = torch.full(
+                (batch_size,),
+                int(audio_tensor.shape[-1]),
+                device=audio_tensor.device,
+                dtype=torch.long,
+            )
+        else:
+            lengths_tensor = torch.as_tensor(audio_lengths, dtype=torch.long, device=audio_tensor.device).reshape(-1)
+            if lengths_tensor.numel() != int(batch_size):
+                raise ValueError(f"Expected {batch_size} decoded audio lengths, got {int(lengths_tensor.numel())}")
+
+        waveform_rows: list[torch.FloatTensor] = []
+        for row_index in range(batch_size):
+            row_length = max(0, int(lengths_tensor[row_index].item()))
+            waveform_rows.append(audio_tensor[row_index, :, :row_length].detach().cpu())
+        return waveform_rows, sample_rate
+
     @staticmethod
     def _resolve_audio_tokenizer_downsample_rate(audio_tokenizer) -> int:
         for holder in (audio_tokenizer, getattr(audio_tokenizer, "config", None)):
@@ -1405,8 +1610,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         )
         return local_outputs.last_hidden_state[:, -1, :]
 
-    @torch.no_grad()
-    def generate(
+    def _iter_generation_events(
         self,
         input_ids: torch.LongTensor,
         attention_mask: Optional[torch.Tensor] = None,
@@ -1422,7 +1626,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         audio_repetition_penalty: float = 1.0,
         use_kv_cache: bool = True,
         return_dict_in_generate: bool = True,
-    ):
+    ) -> Iterator[dict[str, Any]]:
         if input_ids.ndim == 2:
             input_ids = input_ids.unsqueeze(0)
         if input_ids.ndim != 3:
@@ -1442,7 +1646,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         past_key_values = None
         local_dtype = self.local_transformer.ln_f.weight.dtype
 
-        for _ in range(max_new_frames):
+        for step_index in range(max_new_frames):
             generated_audio_history = torch.stack(generated_frames, dim=1) if generated_frames else None
             global_inputs_embeds = self._build_inputs_embeds(current_model_input_ids)
             global_outputs = self.transformer(
@@ -1528,14 +1732,484 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
             else:
                 current_model_input_ids = current_input_ids
 
+            yield {
+                "type": "frame",
+                "step_index": int(step_index),
+                "audio_token_ids": padded_next_frame.detach().clone(),
+                "active_mask": should_continue.detach().clone(),
+                "finished_mask": finished.detach().clone(),
+            }
+
         if generated_frames:
             audio_token_ids = torch.stack(generated_frames, dim=1)
         else:
             audio_token_ids = torch.empty((batch_size, 0, self.config.n_vq), dtype=torch.long, device=input_ids.device)
 
         if not return_dict_in_generate:
-            return audio_token_ids
-        return NanoTTSGenerationOutput(audio_token_ids=audio_token_ids, prompt_input_ids=input_ids)
+            yield {"type": "final", "audio_token_ids": audio_token_ids}
+            return
+        yield {
+            "type": "final",
+            "generation": NanoTTSGenerationOutput(audio_token_ids=audio_token_ids, prompt_input_ids=input_ids),
+        }
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        nq: Optional[int] = None,
+        max_new_frames: int = 300,
+        do_sample: bool = False,
+        text_temperature: float = 1.5,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: float = 1.7,
+        audio_top_p: float = 0.8,
+        audio_top_k: int = 25,
+        audio_repetition_penalty: float = 1.0,
+        use_kv_cache: bool = True,
+        return_dict_in_generate: bool = True,
+    ):
+        final_output: Any = None
+        for event in self._iter_generation_events(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            nq=nq,
+            max_new_frames=max_new_frames,
+            do_sample=do_sample,
+            text_temperature=text_temperature,
+            text_top_p=text_top_p,
+            text_top_k=text_top_k,
+            audio_temperature=audio_temperature,
+            audio_top_p=audio_top_p,
+            audio_top_k=audio_top_k,
+            audio_repetition_penalty=audio_repetition_penalty,
+            use_kv_cache=use_kv_cache,
+            return_dict_in_generate=return_dict_in_generate,
+        ):
+            if event["type"] != "final":
+                continue
+            final_output = event.get("generation", event.get("audio_token_ids"))
+        if final_output is None:
+            raise RuntimeError("Generation finished without producing a final output.")
+        return final_output
+
+    @torch.no_grad()
+    def generate_stream(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        nq: Optional[int] = None,
+        max_new_frames: int = 300,
+        do_sample: bool = False,
+        text_temperature: float = 1.5,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: float = 1.7,
+        audio_top_p: float = 0.8,
+        audio_top_k: int = 25,
+        audio_repetition_penalty: float = 1.0,
+        use_kv_cache: bool = True,
+        return_dict_in_generate: bool = True,
+    ) -> Iterator[dict[str, Any]]:
+        yield from self._iter_generation_events(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            nq=nq,
+            max_new_frames=max_new_frames,
+            do_sample=do_sample,
+            text_temperature=text_temperature,
+            text_top_p=text_top_p,
+            text_top_k=text_top_k,
+            audio_temperature=audio_temperature,
+            audio_top_p=audio_top_p,
+            audio_top_k=audio_top_k,
+            audio_repetition_penalty=audio_repetition_penalty,
+            use_kv_cache=use_kv_cache,
+            return_dict_in_generate=return_dict_in_generate,
+        )
+
+    @torch.no_grad()
+    def inference_stream(
+        self,
+        text: str,
+        output_audio_path: Union[str, Path],
+        mode: str = "continuation",
+        prompt_text: Optional[str] = None,
+        prompt_audio_path: Optional[Union[str, Path]] = None,
+        reference_audio_path: Optional[Union[str, Path]] = None,
+        text_tokenizer=None,
+        text_tokenizer_path: Optional[str] = None,
+        audio_tokenizer=None,
+        audio_tokenizer_type: Optional[str] = None,
+        audio_tokenizer_pretrained_name_or_path: Optional[str] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        nq: Optional[int] = None,
+        max_new_frames: int = 300,
+        do_sample: bool = False,
+        text_temperature: float = 1.5,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: float = 1.7,
+        audio_top_p: float = 0.8,
+        audio_top_k: int = 25,
+        audio_repetition_penalty: float = 1.0,
+        use_kv_cache: bool = True,
+        voice_clone_max_text_tokens: int = DEFAULT_VOICE_CLONE_MAX_TEXT_TOKENS,
+        voice_clone_max_memory_per_sample_gb: float = DEFAULT_VOICE_CLONE_MAX_MEMORY_PER_SAMPLE_GB,
+        tts_max_batch_size: int = 0,
+        codec_max_batch_size: int = 0,
+    ) -> Iterator[dict[str, Any]]:
+        resolved_device = self._resolve_device(device)
+        effective_nq = self._resolve_inference_nq(nq)
+        if next(self.parameters()).device != resolved_device:
+            self.to(resolved_device)
+
+        was_training = self.training
+        self.eval()
+
+        text_tokenizer = self._load_text_tokenizer(
+            text_tokenizer=text_tokenizer,
+            text_tokenizer_path=text_tokenizer_path,
+        )
+        audio_tokenizer = self._load_audio_tokenizer(
+            audio_tokenizer=audio_tokenizer,
+            audio_tokenizer_type=audio_tokenizer_type,
+            audio_tokenizer_pretrained_name_or_path=audio_tokenizer_pretrained_name_or_path,
+            device=resolved_device,
+        )
+
+        target_sample_rate = self._resolve_audio_tokenizer_sample_rate(audio_tokenizer)
+        target_channels = self._resolve_audio_tokenizer_channels(audio_tokenizer)
+        effective_prompt_audio_path = prompt_audio_path or reference_audio_path
+        resolved_mode = self._resolve_inference_mode(
+            mode=mode,
+            has_prompt_text=prompt_text is not None,
+            has_prompt_audio=effective_prompt_audio_path is not None,
+        )
+        if reference_audio_path is not None and prompt_audio_path is None:
+            logging.warning(
+                "reference_audio_path=%s is treated as prompt_audio_path for backward compatibility.",
+                reference_audio_path,
+            )
+
+        prompt_audio_codes = None
+        if effective_prompt_audio_path is not None:
+            waveform, sample_rate = self._load_reference_audio(
+                effective_prompt_audio_path,
+                target_sample_rate,
+                target_channels,
+            )
+            encoded = self._call_audio_encode(
+                audio_tokenizer=audio_tokenizer,
+                waveform=waveform.to(resolved_device),
+                sample_rate=sample_rate,
+            )
+            prompt_audio_codes = self._mask_unused_audio_channels(
+                self._normalize_audio_codes(encoded),
+                nq=effective_nq,
+            ).to(resolved_device)
+
+        if resolved_mode == "voice_clone":
+            split_voice_clone_text_chunks = self._split_text_into_best_sentences(
+                text_tokenizer=text_tokenizer,
+                text=text,
+                max_tokens=voice_clone_max_text_tokens,
+            )
+            voice_clone_text_chunks = split_voice_clone_text_chunks if len(split_voice_clone_text_chunks) > 1 else [text]
+        else:
+            voice_clone_text_chunks = [text]
+
+        if resolved_mode == "voice_clone" and len(voice_clone_text_chunks) > 1:
+            voice_clone_chunk_batch_size, voice_clone_codec_batch_size = self._resolve_effective_voice_clone_batch_sizes(
+                resolved_device=resolved_device,
+                chunk_count=len(voice_clone_text_chunks),
+                max_memory_per_sample_gb=float(voice_clone_max_memory_per_sample_gb),
+                requested_tts_max_batch_size=tts_max_batch_size,
+                requested_codec_max_batch_size=codec_max_batch_size,
+                realtime_streaming=True,
+            )
+        else:
+            voice_clone_chunk_batch_size = 1
+            voice_clone_codec_batch_size = 1
+
+        generated_audio_token_chunks: list[torch.LongTensor] = []
+        emitted_waveform_segments: list[torch.FloatTensor] = []
+        decoded_sample_rate: Optional[int] = None
+        emitted_samples_total = 0
+        first_audio_emitted_at: Optional[float] = None
+        streaming_reset_fn = getattr(audio_tokenizer, "_reset_batch_decode_streaming_state", None)
+
+        try:
+            for batch_start in range(0, len(voice_clone_text_chunks), voice_clone_chunk_batch_size):
+                batch_chunks = voice_clone_text_chunks[batch_start : batch_start + voice_clone_chunk_batch_size]
+                batch_prompt_input_ids: list[torch.LongTensor] = []
+                batch_attention_masks: list[torch.BoolTensor] = []
+                for text_chunk in batch_chunks:
+                    prompt_input_ids, attention_mask = self.build_inference_input_ids(
+                        text=text_chunk,
+                        text_tokenizer=text_tokenizer,
+                        mode=resolved_mode,
+                        prompt_text=prompt_text,
+                        prompt_audio_codes=prompt_audio_codes,
+                        device=resolved_device,
+                    )
+                    batch_prompt_input_ids.append(prompt_input_ids)
+                    batch_attention_masks.append(attention_mask)
+
+                batched_prompt_input_ids, batched_attention_mask = self._left_pad_inference_batch(
+                    input_id_batches=batch_prompt_input_ids,
+                    attention_mask_batches=batch_attention_masks,
+                    device=resolved_device,
+                )
+
+                row_states = [
+                    {
+                        "pending_decode_frames": [],
+                        "decoded_audio_segments": [],
+                        "generation_complete": False,
+                        "pause_emitted": False,
+                    }
+                    for _ in batch_chunks
+                ]
+                batch_emit_index = 0
+                codec_stream_started = False
+
+                if resolved_mode == "continuation" and prompt_audio_codes is not None:
+                    prompt_decode_codes = self._prepare_audio_codes_for_decode(prompt_audio_codes, nq=effective_nq)
+                    _ = audio_tokenizer.batch_decode(
+                        [prompt_decode_codes],
+                        num_quantizers=effective_nq,
+                        streaming=True,
+                        max_batch_size=1,
+                        reset_stream=True,
+                    )
+                    codec_stream_started = True
+
+                def _emit_ready_segments() -> Iterator[dict[str, Any]]:
+                    nonlocal batch_emit_index, emitted_samples_total, first_audio_emitted_at
+                    active_sample_rate = decoded_sample_rate or target_sample_rate
+                    while batch_emit_index < len(batch_chunks):
+                        state = row_states[batch_emit_index]
+                        decoded_segments = state["decoded_audio_segments"]
+                        if decoded_segments:
+                            next_segment = decoded_segments.pop(0)
+                            if next_segment.numel() == 0 or int(next_segment.shape[-1]) <= 0:
+                                continue
+                            emitted_waveform_segments.append(next_segment)
+                            if first_audio_emitted_at is None:
+                                first_audio_emitted_at = time.monotonic()
+                            emitted_samples_total += int(next_segment.shape[-1])
+                            yield {
+                                "type": "audio",
+                                "waveform": next_segment,
+                                "sample_rate": active_sample_rate,
+                                "chunk_index": batch_start + batch_emit_index,
+                                "is_pause": False,
+                                "emitted_audio_seconds": float(emitted_samples_total) / float(active_sample_rate),
+                                "lead_seconds": self._compute_stream_lead_seconds(
+                                    emitted_samples_total=emitted_samples_total,
+                                    sample_rate=active_sample_rate,
+                                    first_audio_emitted_at=first_audio_emitted_at,
+                                ),
+                            }
+                            continue
+
+                        if not state["generation_complete"]:
+                            break
+
+                        if (
+                            resolved_mode == "voice_clone"
+                            and len(voice_clone_text_chunks) > 1
+                            and not state["pause_emitted"]
+                            and (batch_start + batch_emit_index) < len(voice_clone_text_chunks) - 1
+                        ):
+                            state["pause_emitted"] = True
+                            pause_seconds = self._estimate_voice_clone_inter_chunk_pause_seconds(
+                                voice_clone_text_chunks[batch_start + batch_emit_index]
+                            )
+                            pause_samples = max(0, int(round(float(active_sample_rate) * pause_seconds)))
+                            if pause_samples > 0:
+                                silence = torch.zeros((target_channels, pause_samples), dtype=torch.float32)
+                                emitted_waveform_segments.append(silence)
+                                if first_audio_emitted_at is None:
+                                    first_audio_emitted_at = time.monotonic()
+                                emitted_samples_total += int(silence.shape[-1])
+                                yield {
+                                    "type": "audio",
+                                    "waveform": silence,
+                                    "sample_rate": active_sample_rate,
+                                    "chunk_index": batch_start + batch_emit_index,
+                                    "is_pause": True,
+                                    "emitted_audio_seconds": float(emitted_samples_total) / float(active_sample_rate),
+                                    "lead_seconds": self._compute_stream_lead_seconds(
+                                        emitted_samples_total=emitted_samples_total,
+                                        sample_rate=active_sample_rate,
+                                        first_audio_emitted_at=first_audio_emitted_at,
+                                    ),
+                                }
+                            batch_emit_index += 1
+                            continue
+
+                        batch_emit_index += 1
+
+                def _maybe_decode_pending(force: bool) -> Iterator[dict[str, Any]]:
+                    nonlocal codec_stream_started, decoded_sample_rate
+                    pending_counts = [len(state["pending_decode_frames"]) for state in row_states]
+                    total_pending = sum(pending_counts)
+                    if total_pending <= 0:
+                        return
+
+                    active_sample_rate = decoded_sample_rate or target_sample_rate
+                    head_pending = pending_counts[batch_emit_index] if batch_emit_index < len(batch_chunks) else 0
+                    lead_seconds = self._compute_stream_lead_seconds(
+                        emitted_samples_total=emitted_samples_total,
+                        sample_rate=active_sample_rate,
+                        first_audio_emitted_at=first_audio_emitted_at,
+                    )
+                    decode_budget = self._resolve_stream_decode_frame_budget(
+                        emitted_samples_total=emitted_samples_total,
+                        sample_rate=active_sample_rate,
+                        first_audio_emitted_at=first_audio_emitted_at,
+                    )
+
+                    if not force:
+                        should_decode = False
+                        if first_audio_emitted_at is None and head_pending > 0:
+                            should_decode = True
+                        elif head_pending > 0 and lead_seconds < 0.45:
+                            should_decode = True
+                        elif max(pending_counts) >= decode_budget:
+                            should_decode = True
+                        elif lead_seconds < 0.0 and total_pending > 0:
+                            should_decode = True
+                        if not should_decode:
+                            return
+
+                    decode_window = max(pending_counts) if force else max(1, decode_budget)
+                    empty_codes = torch.empty((effective_nq, 0), dtype=torch.long, device=resolved_device)
+                    codes_list: list[torch.Tensor] = []
+                    for state in row_states:
+                        take_count = min(len(state["pending_decode_frames"]), decode_window)
+                        if take_count <= 0:
+                            codes_list.append(empty_codes)
+                            continue
+                        frame_rows = state["pending_decode_frames"][:take_count]
+                        del state["pending_decode_frames"][:take_count]
+                        frame_tensor = torch.cat(frame_rows, dim=0).to(device=resolved_device, dtype=torch.long)
+                        codes_list.append(frame_tensor[:, :effective_nq].transpose(0, 1).contiguous())
+
+                    decode_output = audio_tokenizer.batch_decode(
+                        codes_list,
+                        num_quantizers=effective_nq,
+                        streaming=True,
+                        max_batch_size=(voice_clone_codec_batch_size if not codec_stream_started else None),
+                        reset_stream=not codec_stream_started,
+                    )
+                    codec_stream_started = True
+                    waveform_rows, current_sample_rate = self._extract_batch_waveforms_and_sample_rate(
+                        decode_output,
+                        fallback_sample_rate=target_sample_rate,
+                        batch_size=len(batch_chunks),
+                    )
+                    if decoded_sample_rate is None:
+                        decoded_sample_rate = current_sample_rate
+                    elif decoded_sample_rate != current_sample_rate:
+                        raise ValueError(
+                            f"Decoded sample rates differ across streaming decode calls: {decoded_sample_rate} vs {current_sample_rate}"
+                        )
+
+                    for row_index, waveform_row in enumerate(waveform_rows):
+                        if waveform_row.numel() == 0 or int(waveform_row.shape[-1]) <= 0:
+                            continue
+                        row_states[row_index]["decoded_audio_segments"].append(waveform_row)
+
+                    yield from _emit_ready_segments()
+
+                try:
+                    final_generation = None
+                    for event in self.generate_stream(
+                        input_ids=batched_prompt_input_ids,
+                        attention_mask=batched_attention_mask,
+                        nq=effective_nq,
+                        max_new_frames=max_new_frames,
+                        do_sample=do_sample,
+                        text_temperature=text_temperature,
+                        text_top_p=text_top_p,
+                        text_top_k=text_top_k,
+                        audio_temperature=audio_temperature,
+                        audio_top_p=audio_top_p,
+                        audio_top_k=audio_top_k,
+                        audio_repetition_penalty=audio_repetition_penalty,
+                        use_kv_cache=use_kv_cache,
+                        return_dict_in_generate=True,
+                    ):
+                        if event["type"] == "frame":
+                            frame_audio_token_ids = event["audio_token_ids"]
+                            active_mask = event["active_mask"]
+                            finished_mask = event["finished_mask"]
+                            for row_index in range(len(batch_chunks)):
+                                if bool(active_mask[row_index].item()):
+                                    row_states[row_index]["pending_decode_frames"].append(
+                                        frame_audio_token_ids[row_index : row_index + 1].detach().clone()
+                                    )
+                                if bool(finished_mask[row_index].item()):
+                                    row_states[row_index]["generation_complete"] = True
+                            yield from _maybe_decode_pending(force=False)
+                            continue
+                        final_generation = event.get("generation")
+
+                    if final_generation is None:
+                        raise RuntimeError("Streaming generation finished without a final output.")
+
+                    for state in row_states:
+                        state["generation_complete"] = True
+
+                    yield from _maybe_decode_pending(force=True)
+                    yield from _emit_ready_segments()
+
+                    batched_audio_token_ids = self._mask_unused_audio_channels(final_generation.audio_token_ids, nq=effective_nq)
+                    for sample_index in range(len(batch_chunks)):
+                        generated_audio_token_chunks.append(
+                            self._trim_generated_audio_token_ids(
+                                batched_audio_token_ids[sample_index],
+                                effective_nq=effective_nq,
+                            )
+                        )
+                finally:
+                    if codec_stream_started and callable(streaming_reset_fn):
+                        streaming_reset_fn()
+
+            if generated_audio_token_chunks:
+                audio_token_ids = torch.cat(generated_audio_token_chunks, dim=0)
+            else:
+                audio_token_ids = torch.empty((0, self.config.n_vq), dtype=torch.long, device=resolved_device)
+
+            if emitted_waveform_segments:
+                waveform = torch.cat(emitted_waveform_segments, dim=-1)
+            else:
+                waveform = torch.zeros((target_channels, 0), dtype=torch.float32)
+
+            decoded_sample_rate = decoded_sample_rate or target_sample_rate
+            output_path = Path(output_audio_path)
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            torchaudio.save(str(output_path), waveform, decoded_sample_rate)
+
+            yield {
+                "type": "result",
+                "audio_path": str(output_path),
+                "sample_rate": decoded_sample_rate,
+                "audio_token_ids": audio_token_ids.detach().cpu(),
+                "waveform": waveform,
+                "reference_audio_token_ids": None if prompt_audio_codes is None else prompt_audio_codes.detach().cpu(),
+                "voice_clone_text_chunks": voice_clone_text_chunks,
+                "voice_clone_chunk_batch_size": int(voice_clone_chunk_batch_size),
+                "voice_clone_codec_batch_size": int(voice_clone_codec_batch_size),
+            }
+        finally:
+            if was_training:
+                self.train()
 
     @torch.no_grad()
     def inference(
@@ -1565,6 +2239,8 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         use_kv_cache: bool = True,
         voice_clone_max_text_tokens: int = DEFAULT_VOICE_CLONE_MAX_TEXT_TOKENS,
         voice_clone_max_memory_per_sample_gb: float = DEFAULT_VOICE_CLONE_MAX_MEMORY_PER_SAMPLE_GB,
+        tts_max_batch_size: int = 0,
+        codec_max_batch_size: int = 0,
     ) -> dict[str, Any]:
         resolved_device = self._resolve_device(device)
         effective_nq = self._resolve_inference_nq(nq)
@@ -1631,13 +2307,17 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         decoded_sample_rate: Optional[int] = None
 
         if resolved_mode == "voice_clone" and len(voice_clone_text_chunks) > 1:
-            voice_clone_chunk_batch_size = self._resolve_voice_clone_chunk_batch_size(
+            voice_clone_chunk_batch_size, voice_clone_codec_batch_size = self._resolve_effective_voice_clone_batch_sizes(
                 resolved_device=resolved_device,
                 chunk_count=len(voice_clone_text_chunks),
                 max_memory_per_sample_gb=float(voice_clone_max_memory_per_sample_gb),
+                requested_tts_max_batch_size=tts_max_batch_size,
+                requested_codec_max_batch_size=codec_max_batch_size,
+                realtime_streaming=False,
             )
         else:
             voice_clone_chunk_batch_size = 1
+            voice_clone_codec_batch_size = 1
 
         for batch_start in range(0, len(voice_clone_text_chunks), voice_clone_chunk_batch_size):
             batch_chunks = voice_clone_text_chunks[batch_start : batch_start + voice_clone_chunk_batch_size]
@@ -1677,17 +2357,23 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
                 resolved_device=resolved_device,
             )
 
+            batch_audio_token_chunks: list[torch.LongTensor] = []
             for sample_index in range(len(batch_chunks)):
                 audio_token_ids = self._trim_generated_audio_token_ids(
                     batched_audio_token_ids[sample_index],
                     effective_nq=effective_nq,
                 )
                 generated_audio_token_chunks.append(audio_token_ids)
-
-                if resolved_mode == "voice_clone" and len(voice_clone_text_chunks) > 1:
-                    decoded_waveform, current_sample_rate = self._decode_audio_token_ids_to_waveform(
+                batch_audio_token_chunks.append(audio_token_ids)
+
+            if resolved_mode == "voice_clone" and len(voice_clone_text_chunks) > 1:
+                for codec_batch_start in range(0, len(batch_audio_token_chunks), voice_clone_codec_batch_size):
+                    codec_audio_token_batches = batch_audio_token_chunks[
+                        codec_batch_start : codec_batch_start + voice_clone_codec_batch_size
+                    ]
+                    decoded_waveforms, current_sample_rate = self._decode_audio_token_id_batch_to_waveforms(
                         audio_tokenizer=audio_tokenizer,
-                        audio_token_ids=audio_token_ids,
+                        audio_token_id_batches=codec_audio_token_batches,
                         target_sample_rate=target_sample_rate,
                         effective_nq=effective_nq,
                         resolved_device=resolved_device,
@@ -1698,7 +2384,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
                         raise ValueError(
                             f"Decoded sample rates differ across voice_clone chunks: {decoded_sample_rate} vs {current_sample_rate}"
                         )
-                    decoded_waveform_chunks.append(decoded_waveform)
+                    decoded_waveform_chunks.extend(decoded_waveforms)
 
         if generated_audio_token_chunks:
             audio_token_ids = torch.cat(generated_audio_token_chunks, dim=0)
@@ -1754,4 +2440,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
             "audio_token_ids": audio_token_ids.detach().cpu(),
             "waveform": waveform,
             "reference_audio_token_ids": None if prompt_audio_codes is None else prompt_audio_codes.detach().cpu(),
+            "voice_clone_text_chunks": voice_clone_text_chunks,
+            "voice_clone_chunk_batch_size": int(voice_clone_chunk_batch_size),
+            "voice_clone_codec_batch_size": int(voice_clone_codec_batch_size),
         }

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d19e63fdc6a35f61a7d1c27e06bfacaba7c1ed40ea3c619c86efc64bcd50a496
+oid sha256:24003f2f11ac8a2cbf70514db2d8f1c02fb451aa6b3c0bffc9da09f31cd7caa5
 size 234693095