fix: disable repetition penalty by default for streaming and add windowed penalty

The unbounded repetition penalty was causing sentence repetition by
progressively starving the 2048-token codec vocabulary. Streaming
wrappers now default to repetition_penalty=1.0 (disabled), passed as
an explicit parameter that bypasses _merge_generate_kwargs' 1.05
default. Non-streaming generate_voice_clone() is unaffected.

Also fix decode_padded to pad with -1 instead of 0 (a valid codebook
index), then clamp to >= 0 before decoding, avoiding silent corruption
of the first frames in windowed streaming decode.
This commit is contained in:
Kedara Studios
2026-02-10 10:45:24 +01:00
parent 0cc89f619c
commit c7b9bc6b76
3 changed files with 40 additions and 11 deletions
+14 -7
View File
@@ -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)
@@ -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)
+21 -2
View File
@@ -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