diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py index b34b5b944..99dcb33e4 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py @@ -354,12 +354,11 @@ class _ProjectedTransformer(nn.Module): ) -> None: super().__init__() self.downsample_ratio: int = 1 - # Upstream always materializes a learned Linear here regardless of - # whether input_dimension == d_model — substituting Identity when - # the dims happen to match silently drops a trained projection. - self.in_proj = nn.Linear(input_dimension, d_model, bias=False) + self.in_proj = nn.Linear(input_dimension, d_model, bias=False) if input_dimension != d_model else nn.Identity() self.transformer = _Transformer(d_model=d_model, **kwargs) - self.out_proj = nn.Linear(d_model, output_dimension, bias=False) + self.out_proj = ( + nn.Linear(d_model, output_dimension, bias=False) if output_dimension != d_model else nn.Identity() + ) def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: x = self.in_proj(x.transpose(1, 2)) # (B, D, T) → (B, T, d_model) diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index ea45b9c02..2a1e9a468 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -793,18 +793,23 @@ class MossTTSCodecDecoder(nn.Module): return {f"_codec.{name}" for name, _ in codec.named_parameters()} def _build_codec(self, codec_path: str) -> tuple[Any, nn.Module]: - try: - codec_cfg = MossAudioTokenizerV2Config.from_pretrained(codec_path) - codec = MossAudioTokenizerV2Model(codec_cfg) - logger.info("Using vendored MOSS Audio Tokenizer v2 classes from %s", codec_path) - return codec_cfg, codec - except Exception: - logger.exception( - "Failed to instantiate vendored MOSS Audio Tokenizer v2; falling back to legacy vendored codec." - ) + config_dict, _ = MossAudioTokenizerV2Config.get_config_dict(codec_path) + is_v2 = config_dict.get("number_channels", 1) >= 2 + + if is_v2: + try: + codec_cfg = MossAudioTokenizerV2Config.from_pretrained(codec_path) + codec = MossAudioTokenizerV2Model(codec_cfg) + logger.info("Using vendored MOSS Audio Tokenizer v2 classes from %s", codec_path) + return codec_cfg, codec + except Exception: + logger.exception( + "Failed to instantiate vendored MOSS Audio Tokenizer v2; falling back to legacy vendored codec." + ) codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path) codec = MossAudioTokenizerModel(codec_cfg) + logger.info("Using vendored MOSS Audio Tokenizer v1 classes from %s", codec_path) return codec_cfg, codec def _configure_decoder_cudagraph(self, device: torch.device) -> None: