From 4ec89acf95936134aa4d904f7b20a9fffddc21b4 Mon Sep 17 00:00:00 2001 From: Saleem Djima <69433755+saleemdjima@users.noreply.github.com> Date: Wed, 21 Jan 2026 11:54:53 +0100 Subject: [PATCH] Feat/add chunk size control arg (#56) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This pull request addresses issue #38 **Changes included:** - Added a `max_tokens` parameter to the typer cli for the `generate` command, allowing user to set a default number of tokens per chunk. - Introduced a default variable `MAX_TOKEN_PER_CHUNK` to centralize the default value. - Fixed a tokenization issue by inserting a comma, allowing the model to correctly read each chunk as intended rather than treating one chunk as a single sentence. **Future improvement suggestion :** - Consider enhancing the `split_into_best_sentences` function using regex to split text more accurately without modifying it, which would improve clarity and maintainability. --------- Co-authored-by: Václav Volhejn <8401624+vvolhejn@users.noreply.github.com> --- pocket_tts/data/audio_utils.py | 2 +- pocket_tts/default_parameters.py | 1 + pocket_tts/main.py | 5 +++++ pocket_tts/models/tts_model.py | 12 +++++++++--- 4 files changed, 16 insertions(+), 4 deletions(-) diff --git a/pocket_tts/data/audio_utils.py b/pocket_tts/data/audio_utils.py index 32482e5..bc83435 100644 --- a/pocket_tts/data/audio_utils.py +++ b/pocket_tts/data/audio_utils.py @@ -1,4 +1,4 @@ -"""Various utilities for audio convertion (pcm format, sample rate and channels), +"""Various utilities for audio conversion (pcm format, sample rate and channels), and volume normalization.""" import torch diff --git a/pocket_tts/default_parameters.py b/pocket_tts/default_parameters.py index 390b817..871e0aa 100644 --- a/pocket_tts/default_parameters.py +++ b/pocket_tts/default_parameters.py @@ -5,3 +5,4 @@ DEFAULT_LSD_DECODE_STEPS = 1 DEFAULT_NOISE_CLAMP = None DEFAULT_EOS_THRESHOLD = -4.0 DEFAULT_FRAMES_AFTER_EOS = None +MAX_TOKEN_PER_CHUNK = 50 diff --git a/pocket_tts/main.py b/pocket_tts/main.py index c9711a1..865cb35 100644 --- a/pocket_tts/main.py +++ b/pocket_tts/main.py @@ -22,6 +22,7 @@ from pocket_tts.default_parameters import ( DEFAULT_NOISE_CLAMP, DEFAULT_TEMPERATURE, DEFAULT_VARIANT, + MAX_TOKEN_PER_CHUNK, ) from pocket_tts.models.tts_model import TTSModel from pocket_tts.utils.logging_utils import enable_logging @@ -223,6 +224,9 @@ def generate( str, typer.Option(help="Output path for generated audio") ] = "./tts_output.wav", device: Annotated[str, typer.Option(help="Device to use")] = "cpu", + max_tokens: Annotated[ + int, typer.Option(help="Maximum number of tokens per chunk.") + ] = MAX_TOKEN_PER_CHUNK, ): """Generate speech using Kyutai Pocket TTS.""" if "cuda" in device: @@ -242,6 +246,7 @@ def generate( model_state=model_state_for_voice, text_to_generate=text, frames_after_eos=frames_after_eos, + max_tokens=max_tokens, ) stream_audio_chunks(output_path, audio_chunks, tts_model.config.mimi.sample_rate) diff --git a/pocket_tts/models/tts_model.py b/pocket_tts/models/tts_model.py index 86beb81..8ca017b 100644 --- a/pocket_tts/models/tts_model.py +++ b/pocket_tts/models/tts_model.py @@ -23,6 +23,7 @@ from pocket_tts.default_parameters import ( DEFAULT_NOISE_CLAMP, DEFAULT_TEMPERATURE, DEFAULT_VARIANT, + MAX_TOKEN_PER_CHUNK, ) from pocket_tts.models.flow_lm import FlowLMModel from pocket_tts.models.mimi import MimiModel @@ -363,6 +364,7 @@ class TTSModel(nn.Module): self, model_state: dict, text_to_generate: str, + max_tokens: int = MAX_TOKEN_PER_CHUNK, frames_after_eos: int | None = None, copy_state: bool = True, ) -> torch.Tensor: @@ -405,6 +407,7 @@ class TTSModel(nn.Module): text_to_generate=text_to_generate, frames_after_eos=frames_after_eos, copy_state=copy_state, + max_tokens=max_tokens, ): audio_chunks.append(chunk) return torch.cat(audio_chunks, dim=0) @@ -414,6 +417,7 @@ class TTSModel(nn.Module): self, model_state: dict, text_to_generate: str, + max_tokens: int = MAX_TOKEN_PER_CHUNK, frames_after_eos: int | None = None, copy_state: bool = True, ): @@ -458,7 +462,9 @@ class TTSModel(nn.Module): # by using teacher forcing, but it would be a bit slower. # TODO: add the teacher forcing method for long texts where we use the audio of one chunk # as conditioning for the next chunk. - chunks = split_into_best_sentences(self.flow_lm.conditioner.tokenizer, text_to_generate) + chunks = split_into_best_sentences( + self.flow_lm.conditioner.tokenizer, text_to_generate, max_tokens + ) for chunk in chunks: text_to_generate, frames_after_eos_guess = prepare_text_prompt(chunk) @@ -733,7 +739,7 @@ def prepare_text_prompt(text: str) -> tuple[str, int]: return text, frames_after_eos_guess -def split_into_best_sentences(tokenizer, text_to_generate: str) -> list[str]: +def split_into_best_sentences(tokenizer, text_to_generate: str, max_tokens: int) -> list[str]: text_to_generate, _ = prepare_text_prompt(text_to_generate) text_to_generate = text_to_generate.strip() tokens = tokenizer(text_to_generate) @@ -761,7 +767,7 @@ def split_into_best_sentences(tokenizer, text_to_generate: str) -> list[str]: text = tokenizer.sp.decode(list_of_tokens[start:end]) nb_tokens_and_sentences.append((end - start, text)) - max_nb_tokens_in_a_chunk = 50 + max_nb_tokens_in_a_chunk = max_tokens chunks = [] current_chunk = "" current_nb_of_tokens_in_chunk = 0