diff --git a/qwen_tts/core/models/modeling_qwen3_tts.py b/qwen_tts/core/models/modeling_qwen3_tts.py index 5fabb00..9466d0f 100644 --- a/qwen_tts/core/models/modeling_qwen3_tts.py +++ b/qwen_tts/core/models/modeling_qwen3_tts.py @@ -2625,6 +2625,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) subtalker_temperature: float = 0.9, # Repetition penalty repetition_penalty: float = 1.0, + repetition_penalty_window: int = 100, # Streaming control emit_every_frames: int = 8, decode_window_frames: int = 80, @@ -2653,6 +2654,9 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) top_p: Top-p (nucleus) filtering for sampling temperature: Sampling temperature subtalker_*: Parameters for sub-codebook prediction + repetition_penalty: Penalty factor for previously generated tokens (1.0 = disabled) + repetition_penalty_window: Only penalize tokens from the last N steps (0 = unlimited). + Codec models reuse tokens heavily; unlimited tracking starves the vocabulary. emit_every_frames: Emit PCM chunk every N codec frames decode_window_frames: Window size for decoding (longer = better quality, more latency) overlap_samples: Overlap samples for crossfade between chunks @@ -2729,7 +2733,6 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) token = _sample_next_token(last_logits, temperature, top_k, top_p, suppress_tokens) else: token = torch.argmax(last_logits, dim=-1) - # Debug removed for performance: first token sampled # Extract ref_code for decoder context (if in ICL mode) # This provides stable context from the start, eliminating early voice artifacts @@ -2748,7 +2751,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 + generated_token_ids: list[int] = [token.item()] # 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 @@ -2791,9 +2794,10 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) # Sample next token for first codebook step_logits = step_out.logits[:, -1, :] - # Apply repetition penalty to previously generated tokens + # Apply repetition penalty to recently generated tokens (windowed) if repetition_penalty != 1.0 and len(generated_token_ids) > 0: - prev_ids = torch.tensor(list(set(generated_token_ids)), device=step_logits.device) + recent = generated_token_ids[-repetition_penalty_window:] if repetition_penalty_window > 0 else generated_token_ids + prev_ids = torch.tensor(list(set(recent)), 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) @@ -2948,6 +2952,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) subtalker_temperature: float = 0.9, # Repetition penalty repetition_penalty: float = 1.0, + repetition_penalty_window: int = 100, # Streaming control emit_every_frames: int = 8, decode_window_frames: int = 80, @@ -2980,7 +2985,8 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) top_p: Top-p (nucleus) filtering for sampling temperature: Sampling temperature subtalker_*: Parameters for sub-codebook prediction - repetition_penalty: Penalty to reduce repeated tokens/codes + repetition_penalty: Penalty factor for previously generated tokens (1.0 = disabled) + repetition_penalty_window: Only penalize tokens from the last N steps (0 = unlimited) emit_every_frames: Emit PCM chunk every N codec frames (phase 2) decode_window_frames: Window size for decoding (phase 2) overlap_samples: Overlap samples for crossfade between chunks @@ -3071,7 +3077,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) codes_buffers: list[list[torch.Tensor]] = [[] for _ in range(B)] decoded_tails: list[Optional[np.ndarray]] = [None] * B total_frames_emitted: list[int] = [0] * B - generated_token_ids: list[list[int]] = [[] for _ in range(B)] + generated_token_ids: list[list[int]] = [[token[b].item()] for b in range(B)] finished: list[bool] = [False] * B sr = 24000 # default sample rate, updated on first decode @@ -3124,7 +3130,8 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin) for b in range(B): if finished[b] or len(generated_token_ids[b]) == 0: continue - prev_ids = torch.tensor(list(set(generated_token_ids[b])), device=step_logits.device) + recent = generated_token_ids[b][-repetition_penalty_window:] if repetition_penalty_window > 0 else generated_token_ids[b] + prev_ids = torch.tensor(list(set(recent)), device=step_logits.device) scores = torch.gather(step_logits[b], 0, prev_ids) scores = torch.where(scores > 0, scores / repetition_penalty, scores * repetition_penalty) step_logits[b].scatter_(0, prev_ids, scores) diff --git a/qwen_tts/core/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/qwen_tts/core/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index 67554d9..e9a1816 100644 --- a/qwen_tts/core/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/qwen_tts/core/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -1056,12 +1056,15 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): B, Q, T = codes.shape if T < target_length: - # Pad with zeros on the left - pad = torch.zeros(B, Q, target_length - T, dtype=codes.dtype, device=codes.device) + # Pad with -1 on the left (0 is a valid codebook entry) + pad = torch.full((B, Q, target_length - T), -1, dtype=codes.dtype, device=codes.device) codes_padded = torch.cat([pad, codes], dim=-1) else: codes_padded = codes.contiguous() # Ensure uniform tensor format to avoid recompilation + # Clamp padding to valid range before decoding (codebook indices must be >= 0) + codes_padded = torch.clamp(codes_padded, min=0) + # Run forward (uses compiled path if available) wav = self.forward_optimized(codes_padded) diff --git a/qwen_tts/inference/qwen3_tts_model.py b/qwen_tts/inference/qwen3_tts_model.py index 8c51cf2..c301a3c 100644 --- a/qwen_tts/inference/qwen3_tts_model.py +++ b/qwen_tts/inference/qwen3_tts_model.py @@ -704,6 +704,10 @@ class Qwen3TTSModel: first_chunk_emit_every: int = 0, # 0 = disabled, use emit_every_frames throughout first_chunk_decode_window: int = 48, first_chunk_frames: int = 48, # Switch to stable after this many frames + # Repetition penalty window + repetition_penalty_window: int = 100, + # Repetition penalty (disabled by default for streaming to avoid vocabulary starvation) + repetition_penalty: float = 1.0, **kwargs, ) -> Generator[Tuple[np.ndarray, int], None, None]: """ @@ -728,6 +732,9 @@ class Qwen3TTSModel: first_chunk_emit_every: Emit interval for first chunk phase (0 = disabled). first_chunk_decode_window: Decode window size for first chunk phase. first_chunk_frames: Switch to stable settings after this many frames. + repetition_penalty_window: Only penalize tokens from the last N steps (0 = unlimited). + repetition_penalty: Repetition penalty factor (1.0 = disabled). Disabled by default + for streaming to avoid vocabulary starvation with the small codec vocabulary. **kwargs: Generation parameters (do_sample, top_k, top_p, temperature, etc.) Yields: @@ -787,10 +794,11 @@ class Qwen3TTSModel: # Extract streaming params, filter to only supported ones gen_kwargs = self._merge_generate_kwargs(**kwargs) # Only keep params supported by stream_generate_pcm + # Note: repetition_penalty is passed as an explicit arg, not through gen_kwargs, + # so _merge_generate_kwargs' default (1.05) doesn't override our streaming default (1.0) supported_params = { "do_sample", "top_k", "top_p", "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} @@ -809,6 +817,8 @@ class Qwen3TTSModel: first_chunk_emit_every=first_chunk_emit_every, first_chunk_decode_window=first_chunk_decode_window, first_chunk_frames=first_chunk_frames, + repetition_penalty=repetition_penalty, + repetition_penalty_window=repetition_penalty_window, **gen_kwargs, ): yield chunk, sr @@ -831,6 +841,10 @@ class Qwen3TTSModel: first_chunk_emit_every: int = 0, first_chunk_decode_window: int = 48, first_chunk_frames: int = 48, + # Repetition penalty window + repetition_penalty_window: int = 100, + # Repetition penalty (disabled by default for streaming to avoid vocabulary starvation) + repetition_penalty: float = 1.0, **kwargs, ) -> Generator[Tuple[List[np.ndarray], int], None, None]: """ @@ -853,6 +867,9 @@ class Qwen3TTSModel: first_chunk_emit_every: Emit interval for phase 1 (0 = disabled). first_chunk_decode_window: Decode window for phase 1. first_chunk_frames: Switch to phase 2 after this many frames. + repetition_penalty_window: Only penalize tokens from the last N steps (0 = unlimited). + repetition_penalty: Repetition penalty factor (1.0 = disabled). Disabled by default + for streaming to avoid vocabulary starvation with the small codec vocabulary. **kwargs: Generation parameters (do_sample, top_k, top_p, temperature, etc.) Yields: @@ -915,11 +932,11 @@ class Qwen3TTSModel: ref_ids.append(ref_tok) # Filter to supported generation params + # Note: repetition_penalty is passed as an explicit arg, not through gen_kwargs gen_kwargs = self._merge_generate_kwargs(**kwargs) supported_params = { "do_sample", "top_k", "top_p", "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} @@ -938,6 +955,8 @@ class Qwen3TTSModel: first_chunk_emit_every=first_chunk_emit_every, first_chunk_decode_window=first_chunk_decode_window, first_chunk_frames=first_chunk_frames, + repetition_penalty=repetition_penalty, + repetition_penalty_window=repetition_penalty_window, **gen_kwargs, ): yield chunks_list, sr