import glob import os from collections.abc import Iterable from functools import cached_property import numpy as np import torch import torch.nn as nn from transformers import PretrainedConfig from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( Qwen2_5OmniConfig, Qwen2_5OmniTalkerConfig, Qwen2_5OmniThinkerConfig, ) from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, Qwen2_5OmniThinkerProcessingInfo, ) from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix from vllm.model_executor.models.vision import ( get_llm_pos_ids_for_vision, ) # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( Qwen2_5OmniThinkerDummyInputsBuilder, Qwen2_5OmniThinkerMultiModalProcessor, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges from vllm_omni.platforms import current_omni_platform TALKER_CODEC_EOS_TOKEN_ID = 8294 TALKER_CODEC_BOS_TOKEN_ID = 8293 logger = init_logger(__name__) @MULTIMODAL_REGISTRY.register_processor( Qwen2_5OmniThinkerMultiModalProcessor, info=Qwen2_5OmniThinkerProcessingInfo, dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder, ) class Qwen2_5OmniForConditionalGeneration( nn.Module, SupportsMultiModal, SupportsPP, SupportsMRoPE, Qwen2_5OmniConditionalGenerationMixin, CustomProcessMixin ): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.has_preprocess = False self.have_multimodal_outputs = True config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config multimodal_config = vllm_config.model_config.multimodal_config # keep vllm_config for later submodule init self.vllm_config = vllm_config # Initialize thinker components thinker_config: Qwen2_5OmniThinkerConfig = config.thinker_config self.thinker_config = thinker_config self.multimodal_config = multimodal_config # Initialize talker components talker_config: Qwen2_5OmniTalkerConfig = config.talker_config self.talker_config = talker_config self.model_stage = vllm_config.model_config.model_stage if self.model_stage == "thinker": # Initialize thinker model (multimodal processing) self.thinker = init_vllm_registered_model( vllm_config=vllm_config, hf_config=thinker_config, # Use registry architecture key architectures=["Qwen2_5OmniThinkerModel"], ) self.model = self.thinker self.talker = None self.token2wav = None elif self.model_stage == "talker": multimodal_config.skip_mm_profiling = True self.has_preprocess = True self.set_custom_preprocess(self.talker_preprocess) self.thinker = None # Initialize talker model wrapper (handles projection + LM) with self._mark_language_model(vllm_config=vllm_config): self.talker = init_vllm_registered_model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "talker"), hf_config=talker_config, architectures=["Qwen2_5OmniTalkerModel"], ) self.model = self.talker self.token2wav = None # set suppress start id according to token2wav t2w_token_end_id = getattr( getattr(getattr(config, "token2wav_config", None), "dit_config", None), "num_embeds", None ) if t2w_token_end_id: self.model.set_suppress_start_id(t2w_token_end_id + 1) self.requires_raw_input_tokens = True # Initialize thinker_embedding from config (random weights). # For load_format: dummy this is the final state; # for normal loading, load_weights() overwrites with real weights. self.thinker_embedding = nn.Embedding( self.thinker_config.text_config.vocab_size, self.thinker_config.text_config.hidden_size, ) self._init_special_tokens_embeddings() elif self.model_stage == "code2wav": multimodal_config.skip_mm_profiling = True self.thinker = None self.talker = None # Initialize token2wav (code->mel->wav) like thinker/talker self.token2wav_config = getattr(config, "token2wav_config", None) self.token2wav = None if self.token2wav_config is not None: with self._mark_language_model(vllm_config=vllm_config): self.token2wav = init_vllm_registered_model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "token2wav"), hf_config=self.token2wav_config, architectures=["Qwen2_5OmniToken2WavModel"], ) # voice resources (loaded on demand) self._token2wav_conds: dict[str, torch.Tensor] = {} self._token2wav_ref_mels: dict[str, torch.Tensor] = {} self.model = self.token2wav self.requires_raw_input_tokens = True else: raise ValueError("Invalid model stage") # Set up intermediate tensors self.make_empty_intermediate_tensors = ( (self.thinker.make_empty_intermediate_tensors) if self.model_stage == "thinker" else lambda: None ) def get_language_model(self) -> nn.Module: """Return the language model for MoE detection in upstream load_model. Upstream GPUModelRunner.load_model() calls get_language_model() to find the MoE model inside VLM wrappers. For thinker (which IS MoE), delegate to the child. For talker/code2wav (not MoE, not SupportsMultiModal), return self.model directly to avoid NotImplementedError. """ if hasattr(self.model, "get_language_model"): return self.model.get_language_model() return self.model # -------------------- Device utilities -------------------- @staticmethod def _module_device(module: nn.Module) -> torch.device: try: return next(module.parameters()).device except StopIteration: # No parameters; fall back to buffers or cpu for _, buf in module.named_buffers(recurse=True): return buf.device return torch.device("cpu") def move_submodules_to_devices( self, *, thinker_device: str | torch.device | None = None, talker_device: str | torch.device | None = None, token2wav_device: str | torch.device | None = None, ) -> None: """Optionally move thinker/talker/token2wav to different devices. Example: model.move_submodules_to_devices( thinker_device='cuda:0', talker_device='cuda:1', token2wav_device='cpu', ) """ if thinker_device is not None and self.thinker is not None: self.thinker.to(thinker_device) if talker_device is not None and self.talker is not None: self.talker.to(talker_device) if token2wav_device is not None and self.token2wav is not None: self.token2wav.to(token2wav_device) @cached_property def sampler(self): if hasattr(self.model, "sampler"): return self.model.sampler return Sampler() def embed_input_ids( self, input_ids: torch.Tensor, multimodal_embeddings=None, is_multimodal=None, ) -> torch.Tensor: if self.model_stage == "code2wav": return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size()) if self.model_stage == "talker": return self.model.embed_input_ids(input_ids) return self.model.embed_input_ids( input_ids=input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal ) def embed_multimodal(self, **kwargs): # Delegate to thinker model for multimodal processing return self.model.embed_multimodal(**kwargs) def last_index_of(self, list, value): return len(list) - 1 - list[::-1].index(value) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, generate_audio: bool = True, voice_type: str = "Chelsie", codec: torch.Tensor | None = None, sampling_metadata: SamplingMetadata | None = None, logits_index: int | None = None, sampler=None, additional_information: dict[str, object] | None = None, **kwargs: object, ) -> torch.Tensor | IntermediateTensors | OmniOutput: """ Workflow: 1) Thinker: multimodal understanding → text hidden states. 2) If audio requested and codec not provided, use talker to derive codec. 3) If audio requested (or codec provided), use token2wav to synthesize waveform. 4) Return text hidden states (and audio when applicable). """ if self.model_stage == "thinker": # Normalize to batched inputs if caller provides 1D/2D unbatched tensors # TODO: Remove this hack when NPU supports batched inputs properly added_batch_dim = False if input_ids is not None and input_ids.ndim == 1: input_ids = input_ids.unsqueeze(0) added_batch_dim = True if positions is not None and positions.ndim == 1: positions = positions.unsqueeze(0) added_batch_dim = True if inputs_embeds is not None and inputs_embeds.ndim == 2: inputs_embeds = inputs_embeds.unsqueeze(0) added_batch_dim = True thinker_dev = self._module_device(self.thinker) # if input_ids is None, set it to a zero tensor, in the length of the # same as the embedding seq length if input_ids is None: input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev).unsqueeze( 0 ) # (1, 0) added_batch_dim = True # 1) Thinker (ensure inputs on thinker's device) if input_ids is not None and input_ids.device != thinker_dev: input_ids = input_ids.to(thinker_dev) if positions is not None and positions.device != thinker_dev: positions = positions.to(thinker_dev) if inputs_embeds is not None and inputs_embeds.device != thinker_dev: inputs_embeds = inputs_embeds.to(thinker_dev) if current_omni_platform.is_npu(): # TODO: remove this hack when NPU supports batched inputs properly thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it if positions.ndim == 2 and positions.shape[0] == 3: thinker_positions = positions # MRoPE positions, keep as is else: thinker_positions = positions[0] if positions.ndim > 1 else positions thinker_inputs_embeds = ( inputs_embeds[0] if inputs_embeds is not None and added_batch_dim else inputs_embeds ) else: # Squeeze back if we added batch dim earlier thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it if positions.ndim == 2 and positions.shape[0] == 3: thinker_positions = positions # MRoPE positions, keep as is elif added_batch_dim: thinker_positions = positions[0] else: thinker_positions = positions thinker_inputs_embeds = ( inputs_embeds[0] if inputs_embeds is not None and added_batch_dim else inputs_embeds ) # Run thinker thinker_output = self.thinker( input_ids=thinker_input_ids, positions=thinker_positions, intermediate_tensors=intermediate_tensors, inputs_embeds=thinker_inputs_embeds, **kwargs, ) if isinstance(thinker_output, tuple): embeds, text_hidden_states = thinker_output else: text_hidden_states = thinker_output # Text-only path return OmniOutput( text_hidden_states=(text_hidden_states.reshape(-1, text_hidden_states.shape[-1])), multimodal_outputs=None, ) # 2) Talker (if codec not provided) if self.model_stage == "talker": # mock data for profile if input_ids is None: input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device) self.thinker_reply_part = torch.zeros_like(inputs_embeds) # TODO(Peiqi): temporal hack here to support voice_type. if not hasattr(self, "voice_type"): self.voice_type = voice_type # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it if positions.ndim == 2 and positions.shape[0] == 3: talker_positions = positions # MRoPE positions, keep as is else: talker_positions = positions[0] with torch.inference_mode(): talker_hidden = self.talker( input_ids=input_ids, positions=talker_positions, inputs_embeds=inputs_embeds, ) if sampling_metadata is not None: # the padding token id is set to text model's pad token id, # which do not match with the talker model's word embedding size sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448 return OmniOutput( text_hidden_states=talker_hidden, multimodal_outputs=None, ) if self.model_stage == "code2wav": code = ( input_ids if input_ids is not None else torch.zeros( inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device, ) ) if code.numel() and code[-1] == TALKER_CODEC_EOS_TOKEN_ID: code = code[:-1] if code.numel() and code[0] == TALKER_CODEC_BOS_TOKEN_ID: code = code[1:] audio_tensor = self.generate_audio(code, voice_type) if code.numel() else torch.zeros(0, device=code.device) return OmniOutput(text_hidden_states=None, multimodal_outputs={"model_outputs": audio_tensor}) return OmniOutput( text_hidden_states=torch.cat( [ torch.zeros( [inputs_embeds.shape[0], self.talker.config.hidden_size], dtype=torch.bfloat16, ).to(self._module_device(self.model)), self.talker.thinker_to_talker_proj( self.talker.embed_input_ids( torch.tensor([TALKER_CODEC_BOS_TOKEN_ID, TALKER_CODEC_EOS_TOKEN_ID]) .to(torch.bfloat16) .to(self._module_device(self.model)) ) )[0], ], dim=0, ), multimodal_outputs=None, ) def get_mrope_input_positions( self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec] | None = None, *, hf_config: PretrainedConfig, image_grid_thw: list[list[int]] | torch.Tensor, video_grid_thw: list[list[int]] | torch.Tensor, second_per_grid_ts: list[float] | None = None, context_len: int = 0, seq_len: int | None = None, audio_feature_lengths: torch.Tensor | None = None, use_audio_in_video: bool = False, ) -> tuple[torch.Tensor, int]: """Get mrope input positions and delta value (Qwen2.5-Omni version). Differences from MRotaryEmbedding: 1. Add audio support (and related `audio_feature_lengths`). 2. Add `use_audio_in_video` option to read audio from video inputs. In this case, audio and vision position ids will be split into chunks and interleaved. Example: (V_i are vision position ids, A_i are audio position ids) |V_1 ... V_n|A_1 ... A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|... |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |... """ # TODO(fyabc): refactor and share more code with # _vl_get_input_positions_tensor. thinker_config = hf_config.thinker_config audio_token_id = thinker_config.audio_token_index image_token_id = thinker_config.image_token_index video_token_id = thinker_config.video_token_index audio_start_token_id = thinker_config.audio_start_token_id audio_end_token_id = thinker_config.audio_end_token_id vision_start_token_id = thinker_config.vision_start_token_id vision_end_token_id = thinker_config.vision_end_token_id seconds_per_chunk = thinker_config.seconds_per_chunk spatial_merge_size = thinker_config.vision_config.spatial_merge_size tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25) if isinstance(image_grid_thw, list): image_grid_thw = torch.tensor(image_grid_thw) if isinstance(video_grid_thw, list): video_grid_thw = torch.tensor(video_grid_thw) src_item = input_tokens audio_seqlens = audio_feature_lengths if not second_per_grid_ts: second_per_grid_ts = [1] * video_grid_thw.shape[0] audio_idx = 0 video_idx = 0 image_idx = 0 new_src_item: list[int] = [] llm_pos_ids_list: list[torch.Tensor] = [] idx = 0 while idx < len(src_item): new_src_item_len = len(new_src_item) start_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]: if use_audio_in_video and idx > 0: if src_item[idx] == vision_end_token_id and src_item[idx - 1] == audio_end_token_id: # processing the <|audio_eos|> before <|vision_eos|> start_idx -= 1 elif src_item[idx] == audio_start_token_id and src_item[idx - 1] == vision_start_token_id: # processing the <|audio_bos|> after <|vision_eos|> start_idx -= 1 new_src_item.append(src_item[idx]) llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1) llm_pos_ids_list.append(llm_pos_ids) elif src_item[idx] == audio_token_id: assert audio_seqlens is not None audio_seqlen = audio_seqlens[audio_idx] place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1 new_src_item.extend([audio_token_id] * place_num) llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx llm_pos_ids_list.append(llm_pos_ids) audio_idx += 1 elif src_item[idx] == image_token_id: grid_t = image_grid_thw[image_idx][0] grid_hs = image_grid_thw[:, 1] grid_ws = image_grid_thw[:, 2] t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long() llm_pos_ids = get_llm_pos_ids_for_vision( start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws ) llm_pos_ids_list.append(llm_pos_ids) vision_seqlen = image_grid_thw[image_idx].prod() // (spatial_merge_size**2) new_src_item.extend([image_token_id] * vision_seqlen) image_idx += 1 elif src_item[idx] == video_token_id and not use_audio_in_video: grid_t = video_grid_thw[video_idx][0] grid_hs = video_grid_thw[:, 1] grid_ws = video_grid_thw[:, 2] t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long() llm_pos_ids = get_llm_pos_ids_for_vision( start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws ) llm_pos_ids_list.append(llm_pos_ids) vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2) new_src_item.extend([video_token_id] * vision_seqlen) video_idx += 1 else: # read audio from video assert audio_seqlens is not None audio_seqlen = audio_seqlens[audio_idx] vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2) grid_t = video_grid_thw[video_idx][0] grid_h = video_grid_thw[video_idx][1] grid_w = video_grid_thw[video_idx][2] grid_hs = video_grid_thw[:, 1] grid_ws = video_grid_thw[:, 2] t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk) t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long() t_index_split_chunk = split_list_into_ranges(t_index, t_ntoken_per_chunk) place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2 pure_audio_len = place_num - 2 added_audio_len = 0 audio_llm_pos_ids_list: list[torch.Tensor] = [] for t_chunk in t_index_split_chunk: vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2) new_src_item.extend([video_token_id] * vision_ntoken_per_chunk) vision_llm_pos_ids_list = get_llm_pos_ids_for_vision( start_idx, video_idx, spatial_merge_size, t_chunk, grid_hs, grid_ws, ).split(1, dim=1) llm_pos_ids_list.extend(vision_llm_pos_ids_list) new_src_item.extend(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) * [audio_token_id]) audio_start_idx = ( start_idx if len(audio_llm_pos_ids_list) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1 ) if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0: audio_llm_pos_ids_list = ( torch.arange(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)).expand(3, -1) + audio_start_idx ).split(1, dim=1) else: audio_llm_pos_ids_list = [] added_audio_len += min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) llm_pos_ids_list.extend(audio_llm_pos_ids_list) if added_audio_len < pure_audio_len: new_src_item.extend((pure_audio_len - added_audio_len) * [audio_token_id]) audio_llm_pos_ids_list = ( torch.arange(pure_audio_len - added_audio_len).expand(3, -1) + llm_pos_ids_list[-1].max() + 1 ).split(1, dim=1) llm_pos_ids_list.extend(audio_llm_pos_ids_list) audio_idx += 1 video_idx += 1 # move to the next token idx += len(new_src_item) - new_src_item_len llm_positions = torch.cat(llm_pos_ids_list, dim=1) mrope_position_delta = torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item) llm_positions = llm_positions[:, context_len:seq_len] return llm_positions, mrope_position_delta def generate_audio(self, code, voice_type): token2wav_dev = self._module_device(self.token2wav) if isinstance(code, torch.Tensor): code_tensor = code.to(dtype=torch.long, device=token2wav_dev) else: code_tensor = torch.as_tensor(code, dtype=torch.long, device=token2wav_dev) if code_tensor.ndim == 2 and code_tensor.shape[0] == 1: code_tensor = code_tensor.squeeze(0) audio_tensor = self._codec_to_audio(code_tensor, voice_type) return audio_tensor def _load_talker_embedding( self, ) -> torch.nn.Embedding: return self.talker.language_model.model.embed_tokens def _init_special_tokens_embeddings( self, ): # talker embeddings self.talker_embedding = self._load_talker_embedding() # embed_text_bos_token self.tts_text_spk_token_ids = { # M02: Male voice with standard Mandarin and a slight northern accent "m02": 151870, "Ethan": 151870, # F030: Your anime-styled virtual girlfriend "f030": 151872, "Chelsie": 151872, } self.default_tts_text_spk_type = list(self.tts_text_spk_token_ids.keys())[0] self.tts_text_spk_token_ids["prefix_caching"] = 151870 talker_hf_config = self.talker_config if hasattr(talker_hf_config, "talker_config"): talker_hf_config = talker_hf_config.talker_config self.embed_text_bos_token = self.thinker_embedding( torch.tensor( [talker_hf_config.tts_text_start_token_id], dtype=torch.long, device=self._module_device(self.talker), ) ) self.embed_text_spk_tokens = { key: self.thinker_embedding( torch.tensor( [value], dtype=torch.long, device=self._module_device(self.talker), ) ) for key, value in self.tts_text_spk_token_ids.items() } self.embed_text_eos_token = self.thinker_embedding( torch.tensor( [talker_hf_config.tts_text_end_token_id], dtype=torch.long, device=self._module_device(self.talker), ) ) self.embed_text_pad_token = self.thinker_embedding( torch.tensor( [talker_hf_config.tts_text_pad_token_id], dtype=torch.long, device=self._module_device(self.talker), ) ) self.embed_codec_bos_token = self.talker_embedding( torch.tensor( [talker_hf_config.tts_codec_start_token_id], dtype=torch.long, device=self._module_device(self.talker), ) ) self.embed_codec_eos_token = self.talker_embedding( torch.tensor( [talker_hf_config.tts_codec_end_token_id], dtype=torch.long, device=self._module_device(self.talker), ) ) self.embed_codec_pad_token = self.talker_embedding( torch.tensor( [talker_hf_config.tts_codec_pad_token_id], dtype=torch.long, device=self._module_device(self.talker), ) ) return set(["thinker_embedding.weight", "talker_embedding.weight"]) def _get_embed_text_spk_token(self, voice_type: str): if not hasattr(self, "embed_text_spk_tokens") or voice_type not in self.embed_text_spk_tokens: return self.embed_text_bos_token return self.embed_text_spk_tokens[voice_type] def _get_text_spk_token_id(self, voice_type: str): talker_hf_config = self.talker_config if hasattr(talker_hf_config, "talker_config"): talker_hf_config = talker_hf_config.talker_config if voice_type not in self.tts_text_spk_token_ids: return talker_hf_config.tts_text_start_token_id return self.tts_text_spk_token_ids[voice_type] def talker_preprocess( self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: object, ): # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported. # Rules: # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS] # - Decode segments consist of a single non-special token. # - If additional_information is provided (can be a list split by request or a # concatenated tensor plus a list of shapes), then for each request, reconstruct # the thinker→talker input embeddings for the Prefill segments; # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), # add them; otherwise, keep the original embedding. payload: OmniPayload = info_dict # Ensure we have base embeddings when only ids are provided if input_embeds is None and input_ids is not None: input_embeds = self.talker.embed_input_ids(input_ids) span_len = input_ids.shape[0] if span_len > 1: # prefill return self.thinker_to_talker_process(input_ids, input_embeds, payload) else: # decode return self.thinker_to_talker_decode_one_step(input_ids, input_embeds, payload) def thinker_to_talker_process( self, input_ids: torch.Tensor, input_embeds: torch.Tensor, payload: OmniPayload, ): embed = payload.get("embed", {}) hs = payload.get("hidden_states", {}) ids = payload.get("ids", {}) update_dict = {} prompt_embeds = embed.get("prefill") # Tensor [P,H] thinker_result = hs.get("output") # Tensor [K,H] prompt_token_ids = ids.get("prompt") # list[int] thinker_output_token_ids = ids.get("output") # list[int] if not isinstance(prompt_embeds, torch.Tensor): prompt_embeds = torch.zeros( 0, self.talker.config.hidden_size, dtype=input_embeds.dtype, device=self._module_device(self.model) ) if not isinstance(thinker_result, torch.Tensor): thinker_result = torch.zeros( 0, self.talker.config.hidden_size, dtype=input_embeds.dtype, device=self._module_device(self.model) ) if not isinstance(prompt_token_ids, (list, torch.Tensor)): prompt_token_ids = [] if not isinstance(thinker_output_token_ids, (list, torch.Tensor)): thinker_output_token_ids = [] # TODO(Peiqi): add voice_type support req_input_ids, req_embeds = self._thinker_to_talker_prefill( speaker=self.voice_type, output_prompt_embeds=thinker_result.to(input_embeds.dtype).to(self._module_device(self.model)), output_token_ids=thinker_output_token_ids, thinker_prompt_embeds=prompt_embeds.to(input_embeds.dtype).to(self._module_device(self.model)), prompt_token_ids=prompt_token_ids, ) if thinker_result.ndim == 2 and thinker_result.shape[0] > 0: update_dict.setdefault("embed", {})["thinker_reply"] = thinker_result[1:].detach().to("cpu").contiguous() return req_input_ids, req_embeds, update_dict def _thinker_to_talker_prefill( self, speaker: str, output_prompt_embeds, output_token_ids, thinker_prompt_embeds, prompt_token_ids, ): talker_hf_config = self.talker_config if hasattr(talker_hf_config, "talker_config"): talker_hf_config = talker_hf_config.talker_config # if len(output.outputs[0].token_ids) == 2: # issue request prompt_embeds = torch.cat( [ thinker_prompt_embeds, self._get_embed_text_spk_token(speaker) + self.embed_codec_pad_token, output_prompt_embeds[:1] + self.embed_codec_bos_token, ], dim=0, ) prompt_token_ids_processed = prompt_token_ids + [ talker_hf_config.tts_codec_pad_token_id, output_token_ids[0], ] input_tokens_len = len(prompt_token_ids_processed) # the code below is from model runner in Qwen, may need to further discuss later if input_tokens_len > 2: prompt_token_ids_processed = [self.talker_config.tts_codec_mask_token_id] * (input_tokens_len - 2) + [ self.talker_config.tts_codec_pad_token_id, self.talker_config.tts_codec_start_token_id, ] else: prompt_token_ids_processed = [ self.talker_config.tts_codec_pad_token_id, self.talker_config.tts_codec_start_token_id, ][-input_tokens_len:] if isinstance(prompt_token_ids_processed, list): prompt_token_ids_processed = ( torch.Tensor(prompt_token_ids_processed).to(torch.int64).to(self._module_device(self.talker)) ) return prompt_token_ids_processed, prompt_embeds def thinker_to_talker_decode_one_step(self, input_ids, input_embeds, payload: OmniPayload): embed = payload.get("embed", {}) update_dict = {} # choose step vector in priority order step_vec = None q = embed.get("thinker_reply", None) if isinstance(q, torch.Tensor) and q.numel() > 0: step_vec = q[0:1] new_q = q[1:].detach().to("cpu").contiguous() update_dict.setdefault("embed", {})["thinker_reply"] = new_q else: # B) per-request provided decode vector (optional) dv = embed.get("decode") if isinstance(dv, torch.Tensor) and dv.numel() > 0: step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1) elif ( hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0 ): # C) fallback shared pool step_vec = self.thinker_reply_part[0:1] self.thinker_reply_part = self.thinker_reply_part[1:] if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0: one_id = input_ids[0:1] _, one_embed = self._thinker_to_talker_decode_one_step( output_prompt_embeds=step_vec.to(input_embeds.dtype).to(self._module_device(self.model)), output_token_ids=one_id, ) input_embeds[0] = one_embed[0] return input_ids[0:1], input_embeds[0:1], update_dict def _thinker_to_talker_decode_one_step( self, output_prompt_embeds, output_token_ids, ): processed_output_token_embeds = output_prompt_embeds + self.talker.embed_input_ids( output_token_ids ) # for decode return output_token_ids, processed_output_token_embeds def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, **kwargs: object) -> torch.Tensor | None: # Handle OmniOutput type if isinstance(hidden_states, OmniOutput): hidden_states = hidden_states.text_hidden_states # Use thinker model for logits computation return self.model.compute_logits(hidden_states) def sample( self, logits: torch.Tensor, sampling_metadata: SamplingMetadata, ) -> SamplerOutput | None: # 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__.""" if self.token2wav is None or self.token2wav_config is None: return device = self._module_device(self.token2wav) # optional speaker resources conds = getattr(self.token2wav_config, "conds", None) ref_mels = getattr(self.token2wav_config, "ref_mels", None) if isinstance(conds, dict) and isinstance(ref_mels, dict): self._token2wav_conds = {k: torch.as_tensor(v, device=device) for k, v in conds.items()} self._token2wav_ref_mels = {k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()} # legacy: load from directory if provided model_path = hf_model_folder if isinstance(model_path, str) and os.path.isdir(model_path): spk_pt = os.path.join(model_path, "spk_dict.pt") if os.path.exists(spk_pt): data = torch.load(spk_pt, map_location=device) for key, value in data.items(): self._token2wav_conds[key] = value["cond"].to(device) self._token2wav_ref_mels[key] = value["ref_mel"].to(device) else: # legacy npy inputs for f in sorted(glob.glob(os.path.join(model_path, "inputs", "*spk_emb.npy"))): key = os.path.basename(f).split("_")[0].lower() self._token2wav_conds[key] = torch.as_tensor(np.load(f), device=device) for f in sorted(glob.glob(os.path.join(model_path, "inputs", "*ref_mel.npy"))): key = os.path.basename(f).split("_")[0].lower() 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 voice = voice_type or "default" # Resolve cond / ref_mel if provided cond = None ref_mel = None if voice in self._token2wav_conds and voice in self._token2wav_ref_mels: cond = self._token2wav_conds[voice] ref_mel = self._token2wav_ref_mels[voice] # Fallback: create dummy cond/ref_mel if not provided token2wav_dev = self._module_device(self.token2wav) if cond is None: cond = torch.zeros( (1, self.token2wav_config.dit_config.enc_emb_dim), device=token2wav_dev, dtype=torch.float32, ) if ref_mel is None: ref_mel = torch.zeros( (1, 300, self.token2wav_config.dit_config.mel_dim), device=token2wav_dev, dtype=torch.float32, ) # Ensure codec is (1, T) long tensor on correct device if isinstance(codec_tokens, torch.Tensor): codec = codec_tokens.to(dtype=torch.long, device=token2wav_dev) if codec.ndim == 1: codec = codec.unsqueeze(0) else: codec = torch.as_tensor(codec_tokens, dtype=torch.long, device=token2wav_dev).unsqueeze(0) # Streaming with chunked process and boundary alignment # (rely on token2wav.process_chunk) factor = getattr(self.token2wav.token2wav.factor, "factor", 2) chunk_size = 48 mel_dim = getattr( self.token2wav.token2wav.code2wav_dit_model, "mel_dim", self.token2wav_config.dit_config.mel_dim, ) total_mel = int(codec.shape[1] * factor) steps = 10 # Prepare initial noise for the whole sequence y_all = torch.randn((1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev) logger.info( "Currently, we do not use the chunked process, we only use the " "token2wav.process_chunk for the whole sequence. " "The stream mode will be implemented in the future." ) chunk_ends = [] for i in range(codec.shape[1]): chunk_code_length = i * 2 - 24 finished = i == (codec.shape[1] - 1) if (chunk_code_length > 0 and chunk_code_length % chunk_size == 0) or finished: chunk_ends.append(i) # Number of chunks in mel domain prev_generated = None wav_chunks: list = [] with torch.inference_mode(): for n, i in enumerate([0]): finished = i == codec.shape[1] - 1 _, audio_chunk = self.token2wav.process_chunk( conditioning=cond, reference_mel=ref_mel, codec_all=codec, y_all=y_all, i=n, steps=steps, prev_generated=prev_generated if prev_generated is not None else [], finished=True, ) prev_generated = audio_chunk wav_chunks.append(audio_chunk.detach().cpu().numpy()) if len(wav_chunks) == 0: return torch.zeros(0, device=token2wav_dev) waveform = np.concatenate(wav_chunks) return torch.as_tensor(waveform, device=token2wav_dev) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights for all components of the omni model.""" loaded_weights = set() thinker_weights = [] talker_weights = [] token2wav_weights = [] for k, v in weights: if k.startswith("thinker."): thinker_weights.append((k, v)) elif k.startswith("talker."): talker_weights.append((k, v)) elif k.startswith("token2wav."): token2wav_weights.append((k, v)) else: raise ValueError(f"Unknown weight prefix: {k}") # Load thinker weights if self.thinker: if thinker_weights: thinker_loaded = self.thinker.load_weights(thinker_weights) else: thinker_loaded = set([k for k, v in thinker_weights]) thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker") loaded_weights.update(thinker_loaded) # Load talker weights if talker_weights and self.talker is not None: # Map talker weights to appropriate components if self.thinker is None: thinker_embedding_weights = [w for n, w in thinker_weights if n == "thinker.model.embed_tokens.weight"] if thinker_embedding_weights: self.thinker_embedding = nn.Embedding( thinker_embedding_weights[0].shape[0], thinker_embedding_weights[0].shape[1], ) self.thinker_embedding.weight = nn.Parameter( thinker_embedding_weights[0].to(self._module_device(self.talker)) ) talker_loaded = self.talker.load_weights(talker_weights) talker_loaded = add_prefix_to_loaded_weights(talker_loaded, "talker") loaded_weights.update(talker_loaded) loaded_weights.update(self._init_special_tokens_embeddings()) # Load token2wav weights (if any) if token2wav_weights and self.token2wav is not None: # download weights from huggingface for spk_dict.pt model_path = self.vllm_config.model_config.model download_dir = self.vllm_config.load_config.download_dir if os.path.exists(model_path): hf_model_folder = model_path else: hf_model_folder = download_weights_from_hf_specific( model_path, download_dir, allow_patterns=["*.pt"], ) self._init_token2wav_model(hf_model_folder) t2w_loaded = self.token2wav.load_weights(token2wav_weights, os.path.join(hf_model_folder, "spk_dict.pt")) t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, "token2wav") loaded_weights.update(t2w_loaded) return loaded_weights