mirror of
https://github.com/Nighthawk42/Qwen3-TTS-streaming.git
synced 2026-08-30 08:52:27 +00:00
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:
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user