From f83f18439d2f0d44e132e5ed6a61885b19d19cba Mon Sep 17 00:00:00 2001 From: Kedara Studios Date: Fri, 6 Feb 2026 16:25:05 +0100 Subject: [PATCH] fix: add repetition penalty to streaming and sync upstream finetuning Streaming inference (voice clone) - repetition penalty: Without repetition penalty, the model can fall into a degenerate state where it keeps sampling the same codec tokens over and over. This manifests as: - Looping audio: the same syllable or sound fragment repeats endlessly - Extremely long generation: instead of reaching EOS in ~200-500 frames, it runs for thousands of frames (up to max_frames=10000) - Apparent "slowness": a response that should take ~1s of audio takes 10-30s to generate The fix works by tracking previously generated first-codebook token IDs and penalizing them before sampling: - Tokens with positive logits get divided by repetition_penalty (lowering their probability) - Tokens with negative logits get multiplied by it (pushing them further down) This nudges the model away from re-selecting the same tokens, so it progresses through the text naturally and reaches EOS in a reasonable number of steps rather than looping. Default is 1.0 (disabled) and is exposed through the supported_params whitelist in stream_generate_voice_clone() so it can be set via generate_config or user kwargs. Upstream sync (QwenLM/Qwen3-TTS): - Bump version 0.0.4 -> 0.1.1 to match upstream release. - finetuning/sft_12hz.py: weight sub-talker loss by 0.3 factor to prevent the code predictor gradient from dominating the main talker loss during SFT. - finetuning/sft_12hz.py: remove sub-codebook embedding accumulation loop (codec groups 1-15) from input embeddings, unnecessary and harmful for finetuning convergence (upstream PR #178). - finetuning/README.md: update recommended hyperparameters to batch_size=32, lr=2e-6, num_epochs=10 for more stable training. --- finetuning/README.md | 6 +++--- finetuning/sft_12hz.py | 7 +------ pyproject.toml | 2 +- qwen_tts/core/models/modeling_qwen3_tts.py | 13 +++++++++++++ qwen_tts/inference/qwen3_tts_model.py | 3 ++- 5 files changed, 20 insertions(+), 11 deletions(-) diff --git a/finetuning/README.md b/finetuning/README.md index 7cd48b0..49df275 100644 --- a/finetuning/README.md +++ b/finetuning/README.md @@ -50,9 +50,9 @@ python sft_12hz.py \ --init_model_path Qwen/Qwen3-TTS-12Hz-1.7B-Base \ --output_model_path output \ --train_jsonl train_with_codes.jsonl \ - --batch_size 2 \ - --lr 2e-5 \ - --num_epochs 3 \ + --batch_size 32 \ + --lr 2e-6 \ + --num_epochs 10 \ --speaker_name speaker_test ``` diff --git a/finetuning/sft_12hz.py b/finetuning/sft_12hz.py index 9ed8c48..cb0bd18 100644 --- a/finetuning/sft_12hz.py +++ b/finetuning/sft_12hz.py @@ -92,11 +92,6 @@ def train(): input_embeddings = input_text_embedding + input_codec_embedding - for i in range(1, 16): - codec_i_embedding = model.talker.code_predictor.get_input_embeddings()[i - 1](codec_ids[:, :, i]) - codec_i_embedding = codec_i_embedding * codec_mask.unsqueeze(-1) - input_embeddings = input_embeddings + codec_i_embedding - outputs = model.talker( inputs_embeds=input_embeddings[:, :-1, :], attention_mask=attention_mask[:, :-1], @@ -110,7 +105,7 @@ def train(): sub_talker_logits, sub_talker_loss = model.talker.forward_sub_talker_finetune(talker_codec_ids, talker_hidden_states) - loss = outputs.loss + sub_talker_loss + loss = outputs.loss + 0.3 * sub_talker_loss accelerator.backward(loss) diff --git a/pyproject.toml b/pyproject.toml index b2b180a..d4b6df5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "qwen-tts" -version = "0.0.4" +version = "0.1.1" description = "Qwen-TTS python package" readme = "README.md" requires-python = ">=3.9" diff --git a/qwen_tts/core/models/modeling_qwen3_tts.py b/qwen_tts/core/models/modeling_qwen3_tts.py index 7a632f7..6dc8957 100644 --- a/qwen_tts/core/models/modeling_qwen3_tts.py +++ b/qwen_tts/core/models/modeling_qwen3_tts.py @@ -2598,6 +2598,8 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) subtalker_top_k: int = 50, subtalker_top_p: float = 1.0, subtalker_temperature: float = 0.9, + # Repetition penalty + repetition_penalty: float = 1.0, # Streaming control emit_every_frames: int = 8, decode_window_frames: int = 80, @@ -2711,6 +2713,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) decoded_tail: Optional[np.ndarray] = None frames_since_emit = 0 total_frames_emitted = 0 # Track how many frames we've already emitted audio for + generated_token_ids: list[int] = [] # Track first-codebook tokens for repetition penalty for step_idx in range(max_frames): # Mark step begin for CUDA graphs to avoid tensor overwrite errors @@ -2752,11 +2755,21 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) # Sample next token for first codebook step_logits = step_out.logits[:, -1, :] + + # Apply repetition penalty to previously generated tokens + if repetition_penalty != 1.0 and len(generated_token_ids) > 0: + prev_ids = torch.tensor(list(set(generated_token_ids)), device=step_logits.device) + scores = torch.gather(step_logits[0], 0, prev_ids) + scores = torch.where(scores > 0, scores / repetition_penalty, scores * repetition_penalty) + step_logits[0].scatter_(0, prev_ids, scores) + if do_sample: token = _sample_next_token(step_logits, temperature, top_k, top_p, suppress_tokens) else: token = torch.argmax(step_logits, dim=-1) + generated_token_ids.append(token.item()) + frames_since_emit += 1 # Two-phase streaming: determine current phase settings diff --git a/qwen_tts/inference/qwen3_tts_model.py b/qwen_tts/inference/qwen3_tts_model.py index 75b6d1b..1f4a547 100644 --- a/qwen_tts/inference/qwen3_tts_model.py +++ b/qwen_tts/inference/qwen3_tts_model.py @@ -789,7 +789,8 @@ class Qwen3TTSModel: # Only keep params supported by stream_generate_pcm supported_params = { "do_sample", "top_k", "top_p", "temperature", - "subtalker_dosample", "subtalker_top_k", "subtalker_top_p", "subtalker_temperature" + "subtalker_dosample", "subtalker_top_k", "subtalker_top_p", "subtalker_temperature", + "repetition_penalty" } gen_kwargs = {k: v for k, v in gen_kwargs.items() if k in supported_params}