Nighthawk/Qwen3-TTS-streaming
1
0
Fork 0
mirror of https://github.com/Nighthawk42/Qwen3-TTS-streaming.git synced 2026-08-30 08:52:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: handle multiple EOS tokens for generation termination

Browse Source
This commit is contained in:
Kedara Studios
2026-02-04 19:34:58 +01:00
parent 909e1089be
commit 3640831a47
1 changed files with 30 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
qwen_tts/core/models/modeling_qwen3_tts.py
+30 -7
View File
@@ -2527,6 +2527,17 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
repetition_penalty: float = 1.05,
**kwargs,
):
# Multiple EOS tokens that can terminate generation
eos_ids = {
self.config.talker_config.codec_eos_token_id, # Primary codec EOS
2150, # Codec EOS (model-specific)
2157, # Secondary codec token
151670, # TTS special token
self.config.tts_eos_token_id, # 151673
self.config.im_end_token_id, # 151645
151643, # <|endoftext|>
}
talker_kwargs = {
"max_new_tokens": max_new_tokens,
"min_new_tokens": 2,
@@ -2534,7 +2545,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
"top_k": top_k,
"top_p": top_p,
"temperature": temperature,
"subtalker_dosample": subtalker_dosample,
"subtalker_dosample": subtalker_dosample,
"subtalker_top_k": subtalker_top_k,
"subtalker_top_p": subtalker_top_p,
"subtalker_temperature": subtalker_temperature,
@@ -2545,7 +2556,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
"suppress_tokens": [
i
for i in range(self.config.talker_config.vocab_size - 1024, self.config.talker_config.vocab_size)
if i not in (self.config.talker_config.codec_eos_token_id,)
if i not in eos_ids
],
"output_hidden_states": kwargs.get("output_hidden_states", True),
"return_dict_in_generate": kwargs.get("return_dict_in_generate", True)
@@ -2576,7 +2587,9 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
talker_hidden_states = torch.cat([hid[0][-1][:, -1:] for hid in talker_result.hidden_states], dim=1)[:, :-1]
first_codebook = talker_codes[:, :, 0]
is_stop_token = (first_codebook == self.config.talker_config.codec_eos_token_id)
# Check against all EOS tokens
eos_ids_tensor = torch.tensor(list(eos_ids), device=first_codebook.device, dtype=first_codebook.dtype)
is_stop_token = torch.isin(first_codebook, eos_ids_tensor)
stop_indices = torch.argmax(is_stop_token.int(), dim=1)
has_stop_token = is_stop_token.any(dim=1)
effective_lengths = torch.where(has_stop_token, stop_indices, talker_codes.shape[1])
@@ -2658,13 +2671,23 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
non_streaming_mode=non_streaming_mode,
)
eos_id = self.config.talker_config.codec_eos_token_id
# Multiple EOS tokens that can terminate generation
# Some models may emit different EOS tokens depending on context
eos_ids = {
self.config.talker_config.codec_eos_token_id, # Primary codec EOS
2150, # Codec EOS (model-specific)
2157, # Secondary codec token
151670, # TTS special token
self.config.tts_eos_token_id, # 151673
self.config.im_end_token_id, # 151645
151643, # <|endoftext|>
}
# Build suppress_tokens list (same as in generate())
vocab_size = self.config.talker_config.vocab_size
suppress_tokens = [
i for i in range(vocab_size - 1024, vocab_size)
if i != eos_id
if i not in eos_ids
]
# Mark step begin for CUDA graphs (required for torch.compile with reduce-overhead)
@@ -2750,9 +2773,9 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
# Get codec_ids from hidden_states tuple: (layer_outputs, codec_ids)
codec_ids = step_out.hidden_states[1] # [B, num_code_groups]
# Check for EOS in first codebook ON GPU (avoids CPU sync bottleneck)
# Check for EOS in first codebook
# EOS token is out of range for speech tokenizer, so we must not include it
if codec_ids[0, 0] == eos_id:
if codec_ids[0, 0].item() in eos_ids:
break
# Keep on GPU to avoid CPU<->GPU transfers during decode