[Bugfix][Qwen2.5-Omni] Fix two stale call sites in the speech helpers (#6886)

Signed-off-by: Anai-Guo <antai12232931@outlook.com>
This commit is contained in:
Tai An authored and GitHub committed 2026-09-02 21:41:05 -07:00
1 parent 24c49cf392
commit d1a5a7c178
1 file changed
-51
@@ -862,55 +862,6 @@ class Qwen2_5OmniForConditionalGeneration(
# Use thinker model for sampling
return self.model.sample(logits, sampling_metadata)
def generate_speech(self, text_tokens: torch.Tensor, voice_type: str = "default") -> torch.Tensor:
"""
Generate speech from text tokens using the talker and token2wav models.
This method is kept for backward compatibility and direct speech generation.
Args:
text_tokens: Text tokens from thinker model
voice_type: Voice type for speech generation
Returns:
Audio tensor
"""
# Generate codec tokens using talker model
talker_output = self.talker(input_ids=None, positions=None, inputs_embeds=text_tokens)
# Convert talker output to codec tokens
codec_tokens = self._convert_to_codec_tokens(talker_output)
# Generate audio using token2wav model
return self._codec_to_audio(codec_tokens, voice_type=voice_type)
def _convert_to_codec_tokens(
self, talker_output: torch.Tensor, sampling_metadata: SamplingMetadata
) -> torch.Tensor:
"""
Reference (HF): use the talker's codec head to obtain logits, suppress BOS,
then greedily select the next codec token for the current step.
"""
with torch.inference_mode():
logits = self.talker.compute_logits(talker_output, None)
if logits is None:
return torch.zeros(
(talker_output.size(0), 0),
dtype=torch.long,
device=talker_output.device,
)
# Suppress only codec_bos, consistent with HF generate's
# suppress_tokens behavior
bos_id = None
if hasattr(self, "talker_config") and hasattr(self.talker_config, "tts_codec_start_token_id"):
bos_id = int(getattr(self.talker_config, "tts_codec_start_token_id"))
if bos_id is not None:
logits[..., bos_id] = -1e9
# Take the distribution at the last step and select greedily
next_id = self.talker.sample(logits, sampling_metadata).sampled_token_ids
return next_id.to(dtype=torch.long)
def _init_token2wav_model(self, hf_model_folder):
"""Initialize speaker resources if provided; model is constructed in
__init__."""
@@ -942,8 +893,6 @@ class Qwen2_5OmniForConditionalGeneration(
self._token2wav_ref_mels[key] = torch.as_tensor(np.load(f), device=device)
def _codec_to_audio(self, codec_tokens: torch.Tensor, voice_type: str = "default") -> torch.Tensor | None:
if self.token2wav is None:
self._init_token2wav_model()
if self.token2wav is None:
return None
# Normalize voice type