from collections.abc import Iterable import torch import torch.nn as nn from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeTalkerConfig, ) from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY from vllm.model_executor.models.interfaces import ( SupportsPP, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, maybe_prefix, ) from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp import ( Qwen3OmniMoeTalkerCodePredictor, ) from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( Qwen3MoeLLMForCausalLM, ) from vllm_omni.quantization.component_config import ComponentQuantizationConfig logger = init_logger(__name__) class Qwen3OmniMoeTalkerForConditionalGeneration( nn.Module, SupportsPP, ): """ Qwen3 Omni MoE Talker - Converts text to audio codec codes. The talker is the second stage of Qwen3 Omni MoE's TTS pipeline: 1. Thinker: Generates text response + hidden states 2. Talker: Converts those to 8-layer audio codec codes 3. Code2Wav: Converts codes to waveform ## Key Components: - text_projection: Projects thinker text embeddings → talker dimension - hidden_projection: Projects thinker hidden states → talker dimension - language_model: Main MoE transformer (generates layer 0) - codec_head: Projects to codec vocabulary (layer 0 logits) - code_predictor: Small transformer for layers 1-num_layers-1 """ logger = init_logger(__name__) # Weight mapping from HuggingFace to vLLM naming convention hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ # Main MoE transformer model "talker.model.": "language_model.model.", # Codec head remains separate (outputs audio codes, not text) "talker.codec_head.": "codec_head.", # Code predictor: Now matches HF structure exactly (has .model sub-module) # e.g., "talker.code_predictor.model.codec_embedding.0" → "code_predictor.model.codec_embedding.0" "talker.code_predictor.": "code_predictor.", # Projection layers "talker.text_projection.": "text_projection.", "talker.hidden_projection.": "hidden_projection.", # Fallback: strip talker prefix "talker.": "", } ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config rope_params = getattr(talker_config.text_config, "rope_scaling", None) if rope_params is None: rope_params = getattr(talker_config.text_config, "rope_parameters", None) or {} rope_params = dict(rope_params) # In transformers <5.0.0, rope_theta is a top-level config attribute # (e.g. config.text_config.rope_theta = 1000000.0). # In transformers >=5.0.0 (PR #39847), rope_theta moved inside the # rope_parameters dict (e.g. config.text_config.rope_parameters = # {"rope_theta": 1000000.0, "rope_type": "default"}). # Use setdefault so we never overwrite a value already present. # Precedence: rope_params["rope_theta"] (already set) # > text_config.rope_theta (transformers <5.0.0 top-level attr) # > 1000000 (Qwen3 Omni default) rope_params.setdefault( "rope_theta", getattr(talker_config.text_config, "rope_theta", 1000000), ) talker_config.text_config.rope_parameters = rope_params quant_config = vllm_config.quant_config if isinstance(quant_config, ComponentQuantizationConfig): quant_config = quant_config.resolve("talker") self.quant_config = quant_config self.prefix = prefix self.vllm_config = vllm_config self.config = talker_config self.vocab_size = talker_config.text_config.vocab_size self.router_aux_loss_coef = talker_config.text_config.router_aux_loss_coef self.num_experts = talker_config.text_config.num_experts self.num_experts_per_tok = talker_config.text_config.num_experts_per_tok # thinker projection components for talker self.text_projection = Qwen3OmniMoeTalkerResizeMLP(self.config) self.hidden_projection = Qwen3OmniMoeTalkerResizeMLP(self.config) self.codec_head = nn.Linear(self.config.text_config.hidden_size, self.config.text_config.vocab_size, bias=False) self.rope_deltas = None self.spatial_merge_size = self.config.spatial_merge_size self.vocab_size = self.config.code_predictor_config.vocab_size self.num_code_groups = self.config.code_predictor_config.num_code_groups self.language_model = Qwen3OmniMoeModel( vllm_config=vllm_config, talker_config=self.config, prefix=maybe_prefix(prefix, "language_model"), ) self.code_predictor = Qwen3OmniMoeTalkerCodePredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor") ) def code_predictor_forward( self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor | None = None, *, last_talker_hidden: torch.Tensor | None = None, **_: object, ) -> tuple[torch.Tensor, torch.Tensor]: """Generate full RVQ codes + summed embeddings (single-loop, no KV cache). The code predictor uses re-prefill: each AR step re-forwards the full (short) sequence through the transformer. The returned ``proj_buf`` already contains all codec embeddings at positions 1..G, so summed_embeddings = proj_buf[:, 1:, :].sum(dim=1) — no second loop or re-embedding needed. Returns: result_codes: [batch, num_code_groups, seq_len] summed_embeddings: [batch, seq_len, hidden_size] """ if input_ids is None: raise ValueError("`input_ids` (layer-0 codes) must be provided.") if inputs_embeds is None: raise ValueError("`inputs_embeds` (talker hidden states) must be provided.") if inputs_embeds.ndim == 2: inputs_embeds = inputs_embeds.unsqueeze(0) if input_ids.ndim == 1: input_ids = input_ids.unsqueeze(0) batch_size, seq_len = input_ids.shape device = input_ids.device embed_fn = self.language_model.model.codec_embedding hidden_size = self.config.code_predictor_config.hidden_size result_codes = torch.empty( batch_size, self.num_code_groups, seq_len, dtype=torch.int64, device=device, ) summed_embeddings = torch.empty( batch_size, seq_len, hidden_size, dtype=inputs_embeds.dtype, device=device, ) for pos in range(seq_len): layer0_code = input_ids[:, pos : pos + 1] layer0_embed = embed_fn(layer0_code) pos_all_layers, proj_buf = self.code_predictor( layer0_code, layer0_embed, last_talker_hidden, ) result_codes[:, :, pos : pos + 1] = pos_all_layers # proj_buf layout: [0]=talker_hidden, [1..G]=codec embeds summed_embeddings[:, pos, :] = proj_buf[:, 1:, :].sum(dim=1) return result_codes, summed_embeddings def project_thinker_outputs( self, thinker_embeds: torch.Tensor | None = None, thinker_hidden_states: torch.Tensor | None = None, is_multimodal_mask: torch.Tensor | None = None, ) -> torch.Tensor: """ Project thinker outputs to talker's hidden dimension. The talker has a different hidden size than the thinker, so we need to project the inputs appropriately: - Text embeddings (from thinker's embedding layer) → text_projection - Hidden states (from thinker's last layer, for multimodal) → hidden_projection Args: thinker_embeds: Text embeddings from thinker [batch, seq, thinker_hidden] thinker_hidden_states: Hidden states from thinker's last layer [batch, seq, thinker_hidden] is_multimodal_mask: Boolean mask indicating multimodal positions [batch, seq] Returns: projected_embeds: [batch, seq, talker_hidden] """ if thinker_embeds is None and thinker_hidden_states is None: raise ValueError("Either thinker_embeds or thinker_hidden_states must be provided") # If only embeddings provided, project all as text if thinker_hidden_states is None or is_multimodal_mask is None: return self.text_projection(thinker_embeds) # If only hidden states provided, project all as hidden if thinker_embeds is None: return self.hidden_projection(thinker_hidden_states) # Mixed case: use mask to decide which projection batch_size, seq_len, _ = thinker_embeds.shape output = torch.empty( (batch_size, seq_len, self.config.text_config.hidden_size), device=thinker_embeds.device, dtype=thinker_embeds.dtype, ) # Project multimodal regions using hidden states if is_multimodal_mask.any(): mm_hidden = thinker_hidden_states[is_multimodal_mask] projected_mm = self.hidden_projection(mm_hidden) output[is_multimodal_mask] = projected_mm # Project text regions using embeddings if (~is_multimodal_mask).any(): text_embeds = thinker_embeds[~is_multimodal_mask] projected_text = self.text_projection(text_embeds) output[~is_multimodal_mask] = projected_text return output def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, inputs_embeds: torch.Tensor | None = None, intermediate_tensors: IntermediateTensors | None = None, **kwargs: object, ) -> torch.Tensor | IntermediateTensors: """Forward pass through the talker model.""" if inputs_embeds is None and input_ids is not None: inputs_embeds = self.embed_input_ids(input_ids) input_ids = None talker_hidden_states, _ = self.language_model.model( input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds, **kwargs, ) return talker_hidden_states def compute_logits( self, hidden_states: torch.Tensor, ) -> torch.Tensor | None: """Compute logits for audio codec codes (layer 0 of RVQ). This projects the hidden states to the codec vocabulary space. For full audio generation, layers except 0 would be predicted by the code_predictor after sampling. """ logits = self.codec_head(hidden_states) return logits def make_empty_intermediate_tensors( self, batch_size: int, dtype: torch.dtype, device: torch.device ) -> IntermediateTensors: """Create empty intermediate tensors for pipeline parallelism.""" return self.language_model.make_empty_intermediate_tensors(batch_size, dtype, device) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: """Embed codec input IDs.""" return self.language_model.embed_input_ids(input_ids) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights for the talker model. The weight mapping translates from HuggingFace naming convention to vLLM's internal structure. Code predictor weights are routed to its custom loader for vocab extension support. """ loader = AutoWeightsLoader( self, skip_prefixes=["thinker.", "code2wav."], # "code_predictor."], ) # Don't apply mapper again since we already did it loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) # Log load summary try: total_bytes = 0 for name, param in self.named_parameters(): if param is not None and param.data is not None: total_bytes += param.data.numel() * param.data.element_size() device = next(self.parameters()).device logger.info( "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s", self.__class__.__name__, True, total_bytes / (1024**2), str(device), ) except Exception: logger.error("Error logging model load summary") return loaded class Qwen3OmniMoeTalkerResizeMLP(nn.Module): """ MLP for projecting between thinker and talker hidden dimensions. The thinker and talker have different hidden sizes: - Thinker: config.thinker_hidden_size (e.g., 3584) - Talker: config.text_config.hidden_size (e.g., 2048) This MLP projects from thinker → talker dimension. Two instances are used: - text_projection: For text embeddings from thinker's embedding layer - hidden_projection: For hidden states from thinker's last transformer layer """ def __init__(self, config: Qwen3OmniMoeTalkerConfig): super().__init__() self.linear_fc1 = nn.Linear(config.thinker_hidden_size, config.text_config.intermediate_size, bias=True) self.linear_fc2 = nn.Linear(config.text_config.intermediate_size, config.text_config.hidden_size, bias=True) self.act_fn = _ACTIVATION_REGISTRY[config.text_config.hidden_act] # silu def forward(self, hidden_state): return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state))) class Qwen3OmniMoeModel(Qwen3MoeLLMForCausalLM): """ Qwen3 Omni MoE Talker language model. Extends Qwen3MoeLLMForCausalLM (which already uses FusedMoE with shared-expert support) and replaces the text embedding / LM head with a codec embedding so the talker operates over audio-codec tokens instead of text tokens. """ def __init__(self, vllm_config: VllmConfig, talker_config: Qwen3OmniMoeTalkerConfig, prefix: str): # Create a vllm_config for the talker's text model talker_vllm_config = vllm_config.with_hf_config( talker_config.text_config, architectures=["Qwen3MoeForCausalLM"] ) talker_vllm_config.model_config.hf_text_config = talker_vllm_config.model_config.hf_config super().__init__( vllm_config=talker_vllm_config, prefix=prefix, ) self.config = talker_config self.talker_vllm_config = talker_vllm_config # Remove the inherited LM head so the talker only exposes codec outputs. if hasattr(self, "lm_head"): del self.lm_head # Replace the base embed tokens with codec embedding. if hasattr(self.model, "embed_tokens"): del self.model.embed_tokens # Codec embedding for RVQ code generation self.model.codec_embedding = nn.Embedding( talker_config.text_config.vocab_size, talker_config.text_config.hidden_size, ) def embed_input_ids( self, input_ids: torch.Tensor, ) -> torch.Tensor: """Embed codec input IDs.""" return self.model.codec_embedding(input_ids)