[Bugfix][Model] Fix CoVo-Audio prompt processing and dummy loading (#7909)

Signed-off-by: jeffaa729 <hoiwanglo@gmail.com>

vllm_omni/model_executor/models/covo_audio/covo_audio.py

@@ -21,7 +21,6 @@ from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
     BaseProcessingInfo,
-    ProcessorInputs,
     PromptReplacement,
     PromptUpdateDetails,
 )
@@ -94,20 +93,6 @@ class CovoAudioDummyInputsBuilder(BaseDummyInputsBuilder[CovoAudioProcessingInfo
         dummy_audio = np.zeros((16000 * 30,), dtype=np.float32)
         return {"audio": [(dummy_audio, 16000)] * num_audios}
 
-    def get_dummy_processor_inputs(
-        self,
-        seq_len: int,
-        mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions] | None = None,
-    ) -> ProcessorInputs:
-        dummy_text = self.get_dummy_text(mm_counts)
-        dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
-        dummy_mm_items = self.info.parse_mm_data(dummy_mm_data)
-        return ProcessorInputs(
-            prompt=dummy_text,
-            mm_data_items=dummy_mm_items,
-        )
-
 
 class CovoAudioMultiModalProcessor(OmniMultiModalProcessor[CovoAudioProcessingInfo]):
     def _call_hf_processor(
@@ -186,7 +171,11 @@ class CovoAudioMultiModalProcessor(OmniMultiModalProcessor[CovoAudioProcessingIn
             )
 
         return [
-            PromptReplacement("audio", "<|cAUDIO|>", get_replacement),
+            PromptReplacement(
+                modality="audio",
+                target=[audio_token_id],
+                replacement=get_replacement,
+            ),
         ]
 
 

vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py

@@ -49,7 +49,7 @@ class CovoAudioCode2WavForConditionalGeneration(nn.Module, SupportsPP):
             map_location="cpu",
             weights_only=True,
         )
-        self.decoder.load_state_dict(ckpt)
+        self.decoder.load_state_dict(ckpt, strict=False)
         self.decoder.eval()
         self.infer_config = code2wav_config.inference
         self._decoder_moved = False

vllm_omni/model_executor/models/covo_audio/token2wav.py

@@ -1,9 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 Tencent.
 # token2wav: audio token codes -> waveform (inference only)
 # Pipeline: Token -> Latent (flow matching) -> Waveform (BigVGAN)
 
 import math
-from collections import OrderedDict, namedtuple
+from collections import namedtuple
 
 import numpy as np
 import torch
@@ -803,25 +806,6 @@ class Token2WavDecoder(nn.Module):
         self.upsample_factor = self.token2latent.config.get("upsample_factor", 1)
         self.wav_input_sr = config.get("wav_input_sr", 24000)
 
-        self.trainable_module = ["wavegan", "token2latent"]
-
-    def state_dict(self):
-        param_dict = OrderedDict()
-        for name in self.trainable_module:
-            state = self.get_submodule(name).state_dict(prefix=f"{name}.")
-            param_dict.update(state)
-        return param_dict
-
-    def load_state_dict(self, param_dict):
-        for name in self.trainable_module:
-            module_state = OrderedDict()
-            name_len = len(name)
-            for k, v in param_dict.items():
-                if k.startswith(f"{name}."):
-                    new_k = k[name_len + 1 :]
-                    module_state[new_k] = v
-            self.get_submodule(name).load_state_dict(module_state, strict=False)
-
     @torch.no_grad()
     def preprocess_infer_data(self, data):
         zero_spkr = data.get("zero_spkr", False)