[BugFix] Fix MOSS-TTS codec v1/v2 detection and align vendored tokenizer with upstream (#5635)

Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com>

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)

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: