Merge pull request #5 from rekuenkdr/fix/repetition-penalty-and-finetune-fixes

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 Resolve each training epoch resulting in progressively faster output QwenLM/Qwen3-TTS#178).
    finetuning/README.md: update recommended hyperparameters to batch_size=32, lr=2e-6, num_epochs=10 for more stable training.

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
 ```
 

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)
 

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"

qwen_tts/core/models/modeling_qwen3_tts.py

@@ -2619,6 +2619,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,
@@ -2742,6 +2744,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
@@ -2783,11 +2786,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

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}