[New Model] Add support for tencent/Covo-Audio-Chat (#2293)

Signed-off-by: Dnoob <dxpouo@gmail.com>
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/model_executor/models/covo_audio/config_covo_audio.py

@@ -0,0 +1,72 @@
+# Copyright 2026 Tencent.
+from transformers.configuration_utils import PretrainedConfig
+
+# audio_token_index: token IDs >= this value are audio codes.
+# Matches config.audio_token_index and len(tokenizer) = 151671.
+# Note: the LLM total vocab_size is 168055 (includes audio tokens), which
+# is different from this boundary value.
+COVO_AUDIO_TOKEN_INDEX = 151671
+
+
+class CovoAudioCode2WavConfig(PretrainedConfig):
+    model_type = "covo_audio_code2wav"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.wav_input_sr = 24000
+        self.n_speakers = 1
+        self.n_styles = 1
+
+        self.token2latent = {
+            "upsample_factor": 4,
+            "token_vocab_size": 16384,
+            "token_input_dim": 512,
+            "z_dim": 64,
+            "spkr_embed_dim": 1024,
+            "spkr_mask_ratio": 0.0,
+            "bert_mask_rate0": 0.7,
+            "bert_mask_rate1": 1.0,
+            "random_maskrate": 0.3,
+            "cfg_dropout": 0.2,
+            "sigma": 1e-05,
+            "transformer": {
+                "num_layers": 12,
+                "hidden_size": 1024,
+                "ffn_hidden_size": 4096,
+                "num_heads": 16,
+                "modulation": True,
+                "alibi_bias": False,
+                "rotary_bias": True,
+                "qk_norm": False,
+                "max_position_embeddings": 4096,
+                "attn_dropout": 0.0,
+                "dropout": 0.1,
+            },
+        }
+
+        self.wavegan = {
+            "type": "BigVGANFlowVAE",
+            "upsample_rates": [5, 3, 2, 2, 2, 2],
+            "upsample_kernel_sizes": [10, 6, 4, 4, 4, 4],
+            "upsample_initial_channel": 1536,
+            "resblock": "1",
+            "resblock_kernel_sizes": [3, 7, 11],
+            "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+            "downsample_rates": [2, 2, 2, 2, 3, 5],
+            "downsample_channels": [12, 24, 48, 96, 192, 384, 768],
+            "activation": "snakebeta",
+            "snake_logscale": True,
+            "latent_dim": 64,
+            "use_flow": True,
+            "use_vae": True,
+            "kl_weight": 5,
+            "causal": True,
+            "flow_hidden_channels": 256,
+        }
+
+        self.inference = {
+            "s_steps": 10,
+            "cfg_alpha": 1.0,
+            "dynamic_the": False,
+            "rescale_logits": False,
+        }

vllm_omni/model_executor/models/covo_audio/covo_audio.py

@@ -0,0 +1,426 @@
+# Copyright 2026 Tencent.
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any
+
+import numpy as np
+import torch
+from torch import nn
+from transformers import BatchFeature, WhisperFeatureExtractor
+from vllm.config import VllmConfig
+from vllm.config.multimodal import BaseDummyOptions
+from vllm.inputs import MultiModalDataDict
+from vllm.model_executor.models import SupportsPP
+from vllm.model_executor.models.interfaces import SupportsMultiModal
+from vllm.model_executor.models.utils import init_vllm_registered_model
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalFieldConfig
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    ProcessorInputs,
+    PromptReplacement,
+    PromptUpdateDetails,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
+
+# Max tokens per audio: Whisper 30s → 3000 mel frames → encoder 2x + adapter 8x = 16x → 188
+MAX_AUDIO_TOKENS = 188
+
+
+def _calc_audio_num_tokens(num_samples: int, sample_rate: int) -> int:
+    """Calculate the number of audio tokens for a given audio.
+
+    Matches the original Covo-Audio calc_seq_len: 4 rounds of stride-2
+    downsampling (Whisper encoder 2x + AudioAdapter 3-layer 8x = 16x total).
+    """
+    centiseconds = num_samples * 100 // sample_rate
+    seq_len = centiseconds
+    for _ in range(4):
+        seq_len = (seq_len + 1) // 2
+    return seq_len
+
+
+class CovoAudioProcessingInfo(BaseProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.get_hf_config()
+
+    def get_hf_processor(self, **kwargs):
+        return None
+
+    def get_feature_extractor(self, **kwargs):
+        return WhisperFeatureExtractor(
+            feature_size=128,
+            sampling_rate=16000,
+        )
+
+    def get_data_parser(self) -> MultiModalDataParser:
+        return MultiModalDataParser(
+            target_sr=16000,
+        )
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": None}
+
+    def get_mm_max_tokens_per_item(self, seq_len, mm_counts):
+        return {"audio": MAX_AUDIO_TOKENS}
+
+
+class CovoAudioDummyInputsBuilder(BaseDummyInputsBuilder[CovoAudioProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        num_audios = mm_counts.get("audio", 0)
+        return (" <|begofcAUDIO|><|cAUDIO|><|endofcAUDIO|>" * num_audios).strip()
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+    ) -> MultiModalDataDict:
+        num_audios = mm_counts.get("audio", 0)
+        if num_audios == 0:
+            return {}
+        # 30s at 16kHz = Whisper's max input length, ensures profile_run
+        # allocates memory for the worst-case MAX_AUDIO_TOKENS (188) tokens.
+        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(BaseMultiModalProcessor[CovoAudioProcessingInfo]):
+    def _hf_processor_applies_updates(self, prompt_text, mm_items, hf_processor_mm_kwargs, tokenization_kwargs) -> bool:
+        return False
+
+    def _call_hf_processor(
+        self, prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, Any], tok_kwargs: Mapping[str, object]
+    ) -> BatchFeature:
+        tokenizer = self.info.get_tokenizer()
+        feature_extractor = self.info.get_feature_extractor()
+
+        if isinstance(prompt, str):
+            prompt_ids = tokenizer.encode(prompt)
+        else:
+            prompt_ids = prompt
+
+        audios = mm_data.get("audios", []) or mm_data.get("audio", [])
+        if not audios:
+            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
+
+        raw_audios = []
+        num_tokens_list = []
+        for item in audios:
+            if isinstance(item, tuple):
+                audio_array, _sr = item
+            else:
+                audio_array = item
+            raw_audios.append(audio_array)
+            # MultiModalDataParser resamples to target_sr=16000, so always
+            # compute token count at 16 kHz regardless of the original sr.
+            num_tokens_list.append(_calc_audio_num_tokens(len(audio_array), 16000))
+
+        features = feature_extractor(
+            raw_audios,
+            sampling_rate=16000,
+            return_tensors="pt",
+        )
+
+        return BatchFeature(
+            data={
+                "input_ids": torch.tensor([prompt_ids], dtype=torch.int64),
+                "audio_features": features["input_features"],
+                "audio_num_tokens": torch.tensor(num_tokens_list, dtype=torch.int64),
+            },
+            tensor_type=None,
+        )
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return {
+            "audio_features": MultiModalFieldConfig.batched("audio"),
+            "audio_num_tokens": MultiModalFieldConfig.batched("audio"),
+        }
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs,
+    ) -> Sequence[PromptReplacement]:
+        tokenizer = self.info.get_tokenizer()
+        vocab = tokenizer.get_vocab()
+        audio_token_id = vocab["<|cAUDIO|>"]
+
+        out_mm_data = out_mm_kwargs.get_data()
+
+        def get_replacement(item_idx: int):
+            num_tokens_tensor = out_mm_data.get("audio_num_tokens")
+            if num_tokens_tensor is not None:
+                num_tokens = int(num_tokens_tensor[item_idx])
+            else:
+                num_tokens = MAX_AUDIO_TOKENS
+            return PromptUpdateDetails.select_token_id(
+                [audio_token_id] * num_tokens,
+                embed_token_id=audio_token_id,
+            )
+
+        return [
+            PromptReplacement("audio", "<|cAUDIO|>", get_replacement),
+        ]
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    CovoAudioMultiModalProcessor,
+    info=CovoAudioProcessingInfo,
+    dummy_inputs=CovoAudioDummyInputsBuilder,
+)
+class CovoAudioForConditionalGeneration(
+    nn.Module,
+    SupportsPP,
+    SupportsMultiModal,
+    CustomProcessMixin,
+):
+    requires_raw_input_tokens: bool = True
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality.startswith("audio"):
+            return "<|begofcAUDIO|><|cAUDIO|><|endofcAUDIO|>"
+        return None
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.has_preprocess = False
+        self.have_multimodal_outputs = True
+        self.vllm_config = vllm_config
+
+        self.model_stage = vllm_config.model_config.model_stage
+        if self.model_stage == "fused_thinker_talker":
+            self.has_preprocess = True
+            self.set_custom_preprocess(self.fused_thinker_talker_preprocess)
+            self.fused_thinker_talker = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                hf_config=vllm_config.model_config.hf_config,
+                architectures=["CovoAudioLLMModel"],
+            )
+            self.code2wav = None
+            self.model = self.fused_thinker_talker
+
+        elif self.model_stage == "code2wav":
+            self.code2wav = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                architectures=["CovoAudioCode2WavModel"],
+            )
+            self.fused_thinker_talker = None
+            self.model = self.code2wav
+        else:
+            raise ValueError("Invalid model stage")
+
+        self.make_empty_intermediate_tensors = (
+            self.fused_thinker_talker.make_empty_intermediate_tensors
+            if self.model_stage == "fused_thinker_talker"
+            else lambda *args, **kwargs: None
+        )
+
+    def fused_thinker_talker_preprocess(
+        self,
+        input_ids: torch.Tensor | None,
+        input_embeds: torch.Tensor | None,
+        **info_dict: dict,
+    ):
+        if input_embeds is None and input_ids is not None:
+            input_embeds = self.fused_thinker_talker.embed_input_ids(input_ids)
+        return input_ids, input_embeds, info_dict
+
+    def embed_multimodal(self, **kwargs: object):
+        if self.model_stage == "fused_thinker_talker":
+            return self.fused_thinker_talker.embed_multimodal(**kwargs)
+        return None
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal: bool = False,
+    ) -> torch.Tensor:
+        if self.model_stage == "code2wav":
+            hidden_size = self.vllm_config.model_config.get_hidden_size()
+            return torch.zeros(
+                input_ids.numel(),
+                hidden_size,
+                dtype=self.vllm_config.model_config.dtype,
+                device=input_ids.device,
+            )
+        return self.model.embed_input_ids(input_ids, multimodal_embeddings, is_multimodal=is_multimodal)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata: SamplingMetadata | None = None,
+    ) -> torch.Tensor | None:
+        if isinstance(hidden_states, OmniOutput):
+            hidden_states = hidden_states.text_hidden_states
+        return self.model.compute_logits(hidden_states)
+
+    @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 generate_codes(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        **kwargs: object,
+    ):
+        llm_dev = self._module_device(self.fused_thinker_talker)
+
+        # Fallback: dummy input_ids matching embedding length
+        if input_ids is None:
+            input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=llm_dev).unsqueeze(
+                0
+            )  # (1, seq_len)
+
+        # Thinker (ensure inputs on thinker's device)
+        if input_ids is not None and input_ids.device != llm_dev:
+            input_ids = input_ids.to(llm_dev)
+        if positions is not None and positions.device != llm_dev:
+            positions = positions.to(llm_dev)
+        if inputs_embeds is not None and inputs_embeds.device != llm_dev:
+            inputs_embeds = inputs_embeds.to(llm_dev)
+
+        # Run llm
+        llm_output = self.fused_thinker_talker(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+        # CovoAudio interleaves audio tokens directly in the LLM output
+        # token stream; the stage input processor extracts them from
+        # token_ids, so there is no separate codec output here.
+        return None, llm_output
+
+    def generate_audio(self, code: torch.Tensor):
+        code2wav_dev = self._module_device(self.code2wav)
+
+        if isinstance(code, torch.Tensor):
+            code_tensor = code.to(dtype=torch.long, device=code2wav_dev)
+        else:
+            code_tensor = torch.as_tensor(code, dtype=torch.long, device=code2wav_dev)
+
+        if code_tensor.ndim == 2 and code_tensor.shape[0] == 1:
+            code_tensor = code_tensor.squeeze(0)
+
+        with torch.inference_mode():
+            audio_tensor = self.code2wav(input_ids=code_tensor)
+
+        return audio_tensor
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights for all components of the omni model."""
+        loaded_weights: set[str] = set()
+
+        # code2wav weights are loaded in its __init__ via torch.load;
+        # skip iterating safetensors entirely for this stage.
+        if self.code2wav:
+            c2w_loaded = self.code2wav.load_weights(iter([]))
+            c2w_loaded = add_prefix_to_loaded_weights(c2w_loaded, "code2wav")
+            loaded_weights.update(c2w_loaded)
+            return loaded_weights
+
+        llm_weights = ((k, v) for k, v in weights if k.startswith(("llm", "encoder", "audio_adapter")))
+        if self.fused_thinker_talker:
+            llm_loaded = self.fused_thinker_talker.load_weights(llm_weights)
+            llm_loaded = add_prefix_to_loaded_weights(llm_loaded, "fused_thinker_talker")
+            loaded_weights.update(llm_loaded)
+
+        return loaded_weights
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        sampling_metadata: SamplingMetadata | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors | OmniOutput:
+        """Route forward to the active stage.
+
+        - fused_thinker_talker: run Whisper + LLM → text hidden states
+          (audio tokens are interleaved in the output token stream).
+        - code2wav: run BigVGAN vocoder on audio codes → waveform.
+        """
+        if self.model_stage == "fused_thinker_talker":
+            _, text_hidden_states = self.generate_codes(
+                input_ids=input_ids,
+                positions=positions,
+                inputs_embeds=inputs_embeds,
+                intermediate_tensors=intermediate_tensors,
+                **kwargs,
+            )
+
+            return OmniOutput(
+                text_hidden_states=text_hidden_states.reshape(-1, text_hidden_states.shape[-1]),
+                multimodal_outputs={},
+            )
+
+        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,
+                )
+            )
+
+            # Text-only response: stage processor sends a single sentinel [-1]
+            # token.  Skip BigVGAN inference and return a short silence so the
+            # serving layer produces a valid WAV instead of noise.
+            if code.numel() == 1 and int(code.flatten()[0]) == -1:
+                # 0.1s of silence at 24 kHz
+                return OmniOutput(
+                    text_hidden_states=None,
+                    multimodal_outputs={"model_outputs": torch.zeros(1, 2400, device=code.device)},
+                )
+
+            audio_tensor = self.generate_audio(code)
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={
+                    "model_outputs": audio_tensor.reshape(1, -1) if audio_tensor is not None else audio_tensor
+                },
+            )

vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py

@@ -0,0 +1,126 @@
+# Copyright 2026 Tencent.
+import os
+from collections.abc import Iterable
+
+import numpy as np
+import torch
+from torch import nn
+from vllm.config import VllmConfig
+from vllm.model_executor.models import SupportsPP
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
+
+from .config_covo_audio import CovoAudioCode2WavConfig
+from .token2wav import JsonHParams, Token2WavDecoder
+
+
+class CovoAudioCode2WavForConditionalGeneration(nn.Module, SupportsPP):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        self.vllm_config = vllm_config
+        self.config = vllm_config.model_config.hf_config
+        model_name = vllm_config.model_config.model
+        if os.path.isdir(model_name):
+            model_path = model_name
+        else:
+            from huggingface_hub import snapshot_download
+
+            model_path = snapshot_download(model_name)
+        token2wav_path = os.path.join(model_path, "token2wav")
+
+        code2wav_config = CovoAudioCode2WavConfig()
+        model_config = JsonHParams(
+            **{
+                "token2latent": code2wav_config.token2latent,
+                "wavegan": code2wav_config.wavegan,
+                "wav_input_sr": code2wav_config.wav_input_sr,
+                "global_mean_var": os.path.join(token2wav_path, "global_mean_var.npy"),
+            }
+        )
+
+        self.decoder = Token2WavDecoder(model_config)
+        ckpt = torch.load(
+            os.path.join(token2wav_path, "model.pt"),
+            map_location="cpu",
+            weights_only=True,
+        )
+        self.decoder.load_state_dict(ckpt)
+        self.decoder.eval()
+        self.infer_config = code2wav_config.inference
+        self._decoder_moved = False
+        self._sampler = Sampler()
+
+        # Load bundled speaker prompt files for voice timbre control.
+        prompt_dir = os.path.join(os.path.dirname(__file__), "speaker_prompt")
+        self._prompt_token = torch.from_numpy(np.load(os.path.join(prompt_dir, "prompt_token.npy"))).long()
+        self._prompt_latent = torch.from_numpy(np.load(os.path.join(prompt_dir, "prompt_latent.npy"))).float()
+        self._spkr_embed = torch.from_numpy(np.load(os.path.join(prompt_dir, "speaker_embed.npy"))).float()
+
+    def _ensure_on_device(self, device: torch.device) -> None:
+        """Move decoder and prompt tensors to *device* once."""
+        if self._decoder_moved:
+            return
+        self.decoder = self.decoder.to(device)
+        self._prompt_token = self._prompt_token.to(device)
+        self._prompt_latent = self._prompt_latent.to(device)
+        self._spkr_embed = self._spkr_embed.to(device)
+        self._decoder_moved = True
+
+    def forward(self, input_ids: torch.Tensor, **kwargs):
+        self._ensure_on_device(input_ids.device)
+        data = {
+            "target_token": input_ids.unsqueeze(0),
+            "sample_rate": 24000,
+            "prompt_token": self._prompt_token,
+            "prompt_latent": self._prompt_latent,
+            "spkr_embed": self._spkr_embed,
+        }
+        audio = self.decoder.inference(data, **self.infer_config)
+        return audio
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal: bool = False,
+    ) -> torch.Tensor:
+        hidden_size = self.vllm_config.model_config.get_hidden_size()
+        return torch.zeros(
+            input_ids.numel(),
+            hidden_size,
+            dtype=self.vllm_config.model_config.dtype,
+            device=input_ids.device,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        if hidden_states.numel() == 0:
+            return None
+        batch_size = hidden_states.shape[0] if hidden_states.ndim > 0 else 1
+        vocab_size = self.config.llm_config.vocab_size
+        return torch.zeros(
+            (batch_size, vocab_size),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> SamplerOutput | None:
+        if logits is None or logits.numel() == 0:
+            return None
+        return self._sampler(logits=logits, sampling_metadata=sampling_metadata)
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        # Weights are loaded in __init__ via torch.load
+        # Return all decoder state_dict keys so vLLM knows they are initialized
+        return {f"decoder.{k}" for k in self.decoder.state_dict()}

vllm_omni/model_executor/models/covo_audio/covo_audio_llm.py

@@ -0,0 +1,152 @@
+# Copyright 2026 Tencent.
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+from transformers.models.whisper.modeling_whisper import WhisperEncoder
+from vllm.config import VllmConfig
+from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
+from vllm.sequence import IntermediateTensors
+
+
+class DownsampleLayer(nn.Module):
+    """
+    Downsample layer with 1D convolution and linear layers.
+    """
+
+    def __init__(self, input_dim, output_dim, hidden_dim=2048):
+        super().__init__()
+        self.conv1d = nn.Conv1d(in_channels=input_dim, out_channels=input_dim, kernel_size=3, stride=2, padding=1)
+        self.linear1 = nn.Linear(input_dim, hidden_dim)
+        self.relu1 = nn.ReLU()
+        self.linear2 = nn.Linear(hidden_dim, output_dim)
+        self.relu2 = nn.ReLU()
+
+    def forward(self, x):
+        # x: (B, T, C)
+        x = x.transpose(1, 2)  # -> (B, C, T)
+        x = self.conv1d(x)  # -> (B, C, T // 2)
+        x = x.transpose(1, 2)  # -> (B, T // 2, C)
+        x = self.relu1(x)
+        x = self.linear1(x)  # -> (B, T // 2, hidden_dim)
+        x = self.relu2(x)
+        x = self.linear2(x)  # -> (B, T // 2, output_dim)
+        return x
+
+
+class AudioAdapter(nn.Module):
+    """
+    Audio adapter with downsample layers.
+    """
+
+    def __init__(self, input_dim, output_dim, downsample=8):
+        """
+        Args:
+            input_dim (int): input feature dimension (number of channels)
+            output_dim (int): output feature dimension
+            downsample (int): total downsampling factor, must be a power of 2
+        """
+        super().__init__()
+        assert downsample >= 2 and (downsample & (downsample - 1)) == 0, "downsample must be a power of 2"
+
+        num_layers = (
+            downsample.bit_length() - 1
+        )  # calculate how many downsampling steps are needed to reach the target factor
+
+        layers = []
+        in_dim = input_dim
+        for i in range(num_layers):
+            is_last = i == num_layers - 1
+            out_dim = output_dim if is_last else input_dim
+            layers.append(DownsampleLayer(in_dim, out_dim))
+            in_dim = out_dim
+
+        self.downsample_layers = nn.ModuleList(layers)
+
+    def forward(self, x):
+        """
+        Args:
+            x (Tensor): shape (B, T, C),C=input_dim
+        Returns:
+            Tensor: shape (B, T // downsample, output_dim)
+        """
+        for layer in self.downsample_layers:
+            x = layer(x)
+        return x
+
+
+class CovoAudioLLMForConditionalGeneration(nn.Module, SupportsPP, SupportsMultiModal):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        self.llm = init_vllm_registered_model(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "llm"),
+            architectures=["Qwen2ForCausalLM"],
+        )
+        self.encoder = WhisperEncoder(config.encoder_config)
+        self.audio_adapter = AudioAdapter(
+            config.whisper_feats_dim, config.llm_config.hidden_size, config.adapter_downsample
+        )
+        self.make_empty_intermediate_tensors = self.llm.make_empty_intermediate_tensors
+
+    def get_language_model(self) -> torch.nn.Module:
+        return self.llm
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal=None,
+    ) -> torch.Tensor:
+        return super().embed_input_ids(
+            input_ids,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
+
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+        audio_features = kwargs["audio_features"]
+        audio_num_tokens = kwargs.get("audio_num_tokens")
+        # Match encoder dtype (bfloat16) since audio_features from WhisperFeatureExtractor are float32
+        encoder_dtype = next(self.encoder.parameters()).dtype
+        audio_features = audio_features.to(dtype=encoder_dtype)
+        feats = self.encoder(audio_features).last_hidden_state
+        features = self.audio_adapter(feats)
+        # Truncate each audio to its actual length (Whisper pads all to 30s)
+        # Return a list of per-item tensors since lengths may differ.
+        batch_size = features.shape[0]
+        result = []
+        for i in range(batch_size):
+            if audio_num_tokens is not None:
+                n = int(audio_num_tokens[i])
+                result.append(features[i, :n, :])
+            else:
+                result.append(features[i])
+        return result
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+
+        hidden_state = self.llm(input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds)
+        return hidden_state
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        return self.llm.compute_logits(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/model_executor/models/covo_audio/pipeline.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Covo-Audio pipeline topology (frozen).
+
+Stage 0: fused_thinker_talker — multimodal understanding + interleaved text/audio
+         token generation (AR).
+Stage 1: code2wav              — audio codes → waveform (generation).
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.covo_audio"
+
+COVO_AUDIO_PIPELINE = PipelineConfig(
+    model_type="covo_audio",
+    model_arch="CovoAudioForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="fused_thinker_talker",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="latent",
+            sampling_constraints={
+                "detokenize": True,
+                "stop_token_ids": [151645],
+                "ignore_eos": True,
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="code2wav",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            custom_process_input_func=f"{_PROC}.llm2code2wav",
+            sampling_constraints={"detokenize": False},
+        ),
+    ),
+)

vllm_omni/model_executor/models/covo_audio/prompt_utils.py

@@ -0,0 +1,56 @@
+"""Shared Covo-Audio prompt constants and construction helpers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+COVO_AUDIO_INPUT_PREFIX = "<|begofcAUDIO|><|cAUDIO|><|endofcAUDIO|>"
+
+COVO_AUDIO_SYSTEM_PROMPT = (
+    '你是"小腾"，英文名是"Covo"，由腾讯开发的AI助手。\n'
+    "1、请使用简洁、口语化的语言和用户聊天，"
+    "你的态度积极、耐心，像一位值得信赖的朋友。\n"
+    "2、不要使用列表或编号，避免输出网址、表情符号和复杂的公式。\n"
+    "3、不评价竞争对手，不发表主观政治观点，"
+    "针对色情类、政治类、恐怖类、歧视类、暴力类的用户问题，"
+    "你要妥善应对潜在的安全风险，并给出幽默，情绪安抚以及安全的劝导。\n"
+    "请用文本和音频进行对话，交替生成5个文本token和15个音频token，"
+    "音频部分使用发音人：default_female"
+)
+
+
+def build_covo_audio_chat_prompt(
+    user_content: str,
+    system_prompt: str = COVO_AUDIO_SYSTEM_PROMPT,
+) -> str:
+    """Build a raw chat-template prompt string for Covo-Audio-Chat.
+
+    Used by offline inference and tests where the prompt is passed as a
+    plain string (tokenization happens downstream).
+    """
+    return (
+        f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+        f"<|im_start|>user\n{user_content}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+
+
+def build_covo_audio_prompt_token_ids(
+    tokenizer: Any,
+    user_text: str,
+    system_prompt: str = COVO_AUDIO_SYSTEM_PROMPT,
+) -> list[int]:
+    """Build tokenized prompt IDs for Covo-Audio-Chat.
+
+    Used by the OpenAI-compatible serving layer where prompt_token_ids
+    are passed directly to the engine.
+    """
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_text},
+    ]
+    return tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+    )

vllm_omni/model_executor/models/covo_audio/token2wav.py

@@ -0,0 +1,996 @@
+# Copyright 2026 Tencent.
+# token2wav: audio token codes -> waveform (inference only)
+# Pipeline: Token -> Latent (flow matching) -> Waveform (BigVGAN)
+
+import math
+from collections import OrderedDict, namedtuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch import pow, sin
+from torch.nn import Parameter
+from torch.nn.utils import weight_norm
+
+# --------------- Utilities ---------------
+
+
+class JsonHParams:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, dict):
+                v = JsonHParams(**v)
+            if isinstance(v, str) and v.lower() in ["non", "none", "nil", "null"]:
+                v = None
+            self[k] = v
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def items(self):
+        return self.__dict__.items()
+
+    def values(self):
+        return self.__dict__.values()
+
+    def __len__(self):
+        return len(self.__dict__)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return key in self.__dict__
+
+    def __repr__(self):
+        return self.__dict__.__repr__()
+
+    def get(self, key, default=None):
+        return self.__dict__.get(key, default)
+
+
+def exists(val):
+    return val is not None
+
+
+def eval_decorator(fn):
+    def inner(model, *args, **kwargs):
+        was_training = model.training
+        model.eval()
+        out = fn(model, *args, **kwargs)
+        model.train(was_training)
+        return out
+
+    return inner
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+@torch.autocast(enabled=False, device_type="cuda")
+def apply_rotary_pos_emb(pos, t):
+    return t * pos.cos() + rotate_half(t) * pos.sin()
+
+
+if "sinc" in dir(torch):
+    sinc = torch.sinc
+else:
+
+    def sinc(x: torch.Tensor):
+        return torch.where(
+            x == 0, torch.tensor(1.0, device=x.device, dtype=x.dtype), torch.sin(math.pi * x) / math.pi / x
+        )
+
+
+def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+
+    delta_f = 4 * half_width
+    A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if A > 50.0:
+        beta = 0.1102 * (A - 8.7)
+    elif A >= 21.0:
+        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+    if even:
+        time = torch.arange(-half_size, half_size) + 0.5
+    else:
+        time = torch.arange(kernel_size) - half_size
+    if cutoff == 0:
+        filter_ = torch.zeros_like(time)
+    else:
+        filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+        filter_ /= filter_.sum()
+
+    return filter_.view(1, 1, kernel_size)
+
+
+# --------------- Basic Layers ---------------
+
+
+class Dropout(nn.Module):
+    def __init__(self, p: float = 0.5, inplace: bool = False, force_drop: bool = False, **kwargs):
+        super().__init__()
+        if p < 0.0 or p > 1.0:
+            raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
+        self.p = p
+        self.inplace = inplace
+        self.force_drop = force_drop
+
+    def forward(self, x, **kwargs):
+        return F.dropout(x, p=self.p, training=True if self.force_drop else self.training, inplace=self.inplace)
+
+
+class EmbeddingTable(nn.Embedding):
+    def __init__(self, num_embeddings: int, embedding_dim: int, pad_id=-1, **kwargs):
+        super().__init__(num_embeddings, embedding_dim, **kwargs)
+        nn.init.normal_(self.weight, 0.0, embedding_dim**-0.5)
+        self.pad_id = pad_id
+        self.output_dim = embedding_dim
+
+    def forward(self, x):
+        if self.pad_id is not None:
+            mask = x == self.pad_id
+            x = x.masked_fill(mask, 0)
+        outputs = super().forward(x)
+        if self.pad_id is not None:
+            outputs = outputs.masked_fill(mask.unsqueeze(-1), 0.0)
+        return outputs
+
+
+class Linear(nn.Linear):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        bias: bool = True,
+        w_init_gain: str = "linear",
+        activation=None,
+        **kwargs,
+    ):
+        super().__init__(in_channels, out_channels, bias=bias)
+        self.activation = activation if activation is not None else nn.Identity()
+        self.output_dim = out_channels
+        if w_init_gain is not None:
+            if isinstance(w_init_gain, str):
+                gain = nn.init.calculate_gain(w_init_gain)
+            else:
+                gain = w_init_gain
+            nn.init.xavier_uniform_(self.weight, gain=gain)
+        if bias:
+            nn.init.constant_(self.bias, 0.0)
+
+    def forward(self, x, **kwargs):
+        return self.activation(super().forward(x))
+
+
+class Conv1d(nn.Conv1d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 1,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        padding_mode: str = "zeros",
+        bias: bool = True,
+        padding=None,
+        causal: bool = False,
+        bn: bool = False,
+        activation=None,
+        w_init_gain=None,
+        input_transpose: bool = False,
+        **kwargs,
+    ):
+        self.causal = causal
+        if padding is None:
+            if causal:
+                padding = 0
+                self.left_padding = dilation * (kernel_size - 1)
+            else:
+                padding = get_padding(kernel_size, dilation)
+
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            padding_mode=padding_mode,
+            bias=bias,
+        )
+
+        self.in_channels = in_channels
+        self.transpose = input_transpose
+        self.bn = nn.BatchNorm1d(out_channels) if bn else nn.Identity()
+        self.activation = activation if activation is not None else nn.Identity()
+        if w_init_gain is not None:
+            nn.init.xavier_uniform_(self.weight, gain=nn.init.calculate_gain(w_init_gain))
+
+    def forward(self, x):
+        if self.transpose or x.size(1) != self.in_channels:
+            assert x.size(2) == self.in_channels
+            x = x.transpose(1, 2)
+            self.transpose = True
+
+        if self.causal:
+            x = F.pad(x.unsqueeze(2), (self.left_padding, 0, 0, 0)).squeeze(2)
+
+        outputs = self.activation(self.bn(super().forward(x)))
+        return outputs.transpose(1, 2) if self.transpose else outputs
+
+
+class ConvTranspose1d(nn.ConvTranspose1d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        output_padding: int = 0,
+        groups: int = 1,
+        bias: bool = True,
+        dilation: int = 1,
+        padding=None,
+        padding_mode: str = "zeros",
+        causal: bool = False,
+        input_transpose: bool = False,
+        **kwargs,
+    ):
+        if padding is None:
+            padding = 0 if causal else (kernel_size - stride) // 2
+        if causal:
+            assert padding == 0, "padding is not allowed in causal ConvTranspose1d."
+            assert kernel_size == 2 * stride, "kernel_size must be equal to 2*stride in Causal ConvTranspose1d."
+
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            output_padding=output_padding,
+            groups=groups,
+            bias=bias,
+            dilation=dilation,
+            padding_mode=padding_mode,
+        )
+
+        self.causal = causal
+        self.stride = stride
+        self.transpose = input_transpose
+
+    def forward(self, x):
+        if self.transpose or x.size(1) != self.in_channels:
+            assert x.size(2) == self.in_channels
+            x = x.transpose(1, 2)
+            self.transpose = True
+
+        x = super().forward(x)
+        if self.causal:
+            x = x[:, :, : -self.stride]
+        return x.transpose(1, 2) if self.transpose else x
+
+
+class ConvPositionEmbed(nn.Module):
+    def __init__(self, hidden_size: int, kernel_size: int = 31, groups: int = 16):
+        super().__init__()
+        self.conv = Conv1d(
+            hidden_size, hidden_size, kernel_size, groups=groups, input_transpose=True, activation=nn.GELU()
+        )
+
+    def forward(self, x, mask=None):
+        if mask is not None:
+            mask = mask[..., None]
+            x = x.masked_fill(~mask, 0.0)
+        x = self.conv(x)
+        if mask is not None:
+            x = x.masked_fill(~mask, 0.0)
+        return x
+
+
+# --------------- Anti-aliasing (BigVGAN) ---------------
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff=0.5,
+        half_width=0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ):
+        super().__init__()
+        if cutoff < -0.0:
+            raise ValueError("Minimum cutoff must be larger than zero.")
+        if cutoff > 0.5:
+            raise ValueError("A cutoff above 0.5 does not make sense.")
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+        self.register_buffer("filter", filter)
+
+    def forward(self, x):
+        _, C, _ = x.shape
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+        return out
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.stride = ratio
+        self.pad = self.kernel_size // ratio - 1
+        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
+        self.register_buffer("filter", filter)
+
+    def forward(self, x):
+        _, C, _ = x.shape
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+        x = x[..., self.pad_left : -self.pad_right]
+        return x
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=self.kernel_size
+        )
+
+    def forward(self, x):
+        return self.lowpass(x)
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self, activation, up_ratio: int = 2, down_ratio: int = 2, up_kernel_size: int = 12, down_kernel_size: int = 12
+    ):
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.act(x)
+        x = self.downsample(x)
+        return x
+
+
+# --------------- Activations ---------------
+
+
+class SnakeBeta(nn.Module):
+    def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
+        super().__init__()
+        self.in_features = in_features
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+            self.beta = Parameter(torch.zeros(in_features) * alpha)
+        else:
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+            self.beta = Parameter(torch.ones(in_features) * alpha)
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+        return x
+
+
+# --------------- Attention ---------------
+
+AttentionConfig = namedtuple("AttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"])
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, theta=50000):
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+    @property
+    def device(self):
+        return self.inv_freq.device
+
+    @torch.autocast(enabled=False, device_type="cuda")
+    def forward(self, t):
+        if not torch.is_tensor(t):
+            t = torch.arange(t, device=self.device)
+        t = t.type_as(self.inv_freq)
+        freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
+        freqs = torch.cat((freqs, freqs), dim=-1)
+        return freqs
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.0,
+        dropout: float = 0.0,
+        max_position_embeddings: int = 4096,
+        norm_layer: nn.Module = nn.LayerNorm,
+        rotary_bias: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        assert hidden_size % num_heads == 0, "hidden_size should be divisible by num_heads"
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.scale = self.head_dim**-0.5
+        self.max_position_embeddings = max_position_embeddings
+        self.rotary_bias = rotary_bias
+
+        self.q_proj = Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.k_proj = Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.v_proj = Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+
+        self.attn_drop = Dropout(attn_drop)
+        self.o_proj = Linear(hidden_size, hidden_size)
+        self.o_dropout = Dropout(dropout)
+
+        self.cpu_config = AttentionConfig(True, True, True)
+        device_properties = torch.cuda.get_device_properties(torch.device("cuda"))
+        if device_properties.major == 8 and device_properties.minor == 0:
+            self.cuda_config = AttentionConfig(True, True, True)
+        else:
+            self.cuda_config = AttentionConfig(False, True, True)
+
+        if self.rotary_bias:
+            self.rotary = RotaryEmbedding(self.head_dim)
+
+    def forward(self, q, k=None, v=None, mask=None):
+        k = k or q
+        v = v or q
+        B, L, C = q.shape
+        B, S, C = v.shape
+        if mask is not None:
+            if mask.ndim == 2:
+                assert L == S
+                mask = rearrange(mask, "b j -> b 1 1 j")
+                mask = mask.expand(-1, self.num_heads, L, -1)
+            elif mask.ndim == 3:
+                assert mask.size(1) == L and mask.size(2) == S
+                mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
+
+        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
+        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
+        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)
+        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        config = self.cuda_config if q.is_cuda else self.cpu_config
+        attn_bias = torch.zeros(B, self.num_heads, L, S, dtype=q.dtype, device=q.device)
+
+        if self.rotary_bias:
+            if L == S:
+                rotary_emb = self.rotary(L)
+                q, k = map(lambda x: apply_rotary_pos_emb(rotary_emb, x), (q, k))
+            else:
+                q_rotary_emb = self.rotary(L)
+                k_rotary_emb = self.rotary(S)
+                q = apply_rotary_pos_emb(q_rotary_emb, q)
+                k = apply_rotary_pos_emb(k_rotary_emb, k)
+
+        if mask is not None:
+            attn_bias.masked_fill_(mask.logical_not(), float("-inf"))
+
+        with torch.backends.cuda.sdp_kernel(**config._asdict()):
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=attn_bias, dropout_p=self.attn_drop.p if self.training else 0.0
+            )
+
+        out = rearrange(out, "b h n d -> b n (h d)")
+        out = self.o_dropout(self.o_proj(out))
+        return out
+
+
+# --------------- DiT Modules ---------------
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
+            device=t.device
+        )
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, output_size):
+        super().__init__()
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = Linear(hidden_size, output_size, bias=True)
+
+    def forward(self, x, c, mask=None):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+# --------------- Transformer ---------------
+
+
+class Mlp(nn.Module):
+    def __init__(self, hidden_size, ffn_hidden_size=4096, act_layer=nn.GELU, dropout=0.0, **kwargs):
+        super().__init__()
+        self.fc1 = Linear(hidden_size, ffn_hidden_size)
+        self.act = act_layer()
+        self.fc2 = Linear(ffn_hidden_size, hidden_size)
+        self.drop = Dropout(dropout)
+
+    def forward(self, x, **kwargs):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        attention: nn.Module,
+        ffn: nn.Module,
+        hidden_size: int = 1024,
+        modulation: bool = False,
+        eps: float = 1e-6,
+        **kwargs,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps)
+        self.attn = attention
+        self.ffn = ffn
+        self.modulation = modulation
+        if modulation:
+            self.modulation_layer = nn.Sequential(nn.SiLU(), Linear(hidden_size, 6 * hidden_size, bias=True))
+            nn.init.constant_(self.modulation_layer[-1].weight, 0.0)
+            nn.init.constant_(self.modulation_layer[-1].bias, 0.0)
+
+    def forward(self, x, condition=None, mask=None):
+        if condition is None:
+            assert not self.modulation, "Without global condition, must set modulation to False"
+        else:
+            assert self.modulation, "With global condition, must set modulation to True"
+            shift_attn, scale_attn, gate_attn, shift_ffn, scale_ffn, gate_ffn = self.modulation_layer(condition).chunk(
+                6, dim=1
+            )
+
+        if condition is not None:
+            x = x + gate_attn.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_attn, scale_attn), mask=mask)
+        else:
+            x = x + self.attn(self.norm1(x), mask=mask)
+
+        if condition is not None:
+            x = x + gate_ffn.unsqueeze(1) * self.ffn(modulate(self.norm2(x), shift_ffn, scale_ffn))
+        else:
+            x = x + self.ffn(self.norm2(x), mask=mask)
+        return x
+
+
+# --------------- Flow Matching: Token -> Latent ---------------
+
+
+class Token2latentFlowMatching(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.model_dim = config.transformer.hidden_size
+        self.token_input_dim = config.get("token_input_dim", self.model_dim)
+        self.target_dim = config.z_dim
+
+        self.spkr_embed_dim = config.get("spkr_embed_dim", 512)
+        self.cond_proj = Linear(self.model_dim + self.spkr_embed_dim, self.model_dim)
+        self.token_pad_id = -1
+        if config.upsample_factor > 1:
+            self.token_proj = nn.Sequential(
+                Linear(self.token_input_dim, self.model_dim, bias=True),
+                ConvTranspose1d(
+                    self.model_dim,
+                    self.model_dim,
+                    stride=config.upsample_factor,
+                    kernel_size=config.upsample_factor * 2,
+                ),
+            )
+        else:
+            self.token_proj = Linear(self.token_input_dim, self.model_dim)
+
+        self.transformer_input_proj = Linear(self.model_dim + self.target_dim * 2, self.model_dim)
+        self.time_embedder = TimestepEmbedder(self.model_dim)
+        self.conv_embed = ConvPositionEmbed(hidden_size=self.model_dim, kernel_size=31, groups=16)
+        self.blocks = nn.ModuleList()
+        for _ in range(config.transformer.num_layers):
+            attn_block = MultiHeadAttention(**config.transformer)
+            ffn_block = Mlp(act_layer=lambda: nn.GELU(approximate="tanh"), **config.transformer)
+            self.blocks.append(TransformerBlock(attn_block, ffn_block, **config.transformer))
+        self.output_layer = FinalLayer(config.transformer.hidden_size, self.target_dim)
+
+    def cond_mask_spkr_embed(self, x, spkr_embed):
+        b, device = x.size(0), x.device
+        if not exists(spkr_embed):
+            spkr_embed = torch.zeros(b, self.spkr_embed_dim, device=device, dtype=x.dtype)
+        return spkr_embed
+
+    def vectorfield_forward(self, inputs, times, self_attn_mask, g_cond=None):
+        t = self.time_embedder(times)
+        if g_cond is None:
+            c = t
+        else:
+            cond_inp = torch.cat([t, g_cond], dim=-1)
+            c = self.cond_proj(cond_inp)
+
+        ut = self.transformer_input_proj(inputs)
+        ut = self.conv_embed(ut, mask=self_attn_mask) + ut
+
+        for block in self.blocks:
+            ut = block(ut, c, mask=self_attn_mask)
+        ut = self.output_layer(ut, c)
+        return ut
+
+    @eval_decorator
+    @torch.no_grad()
+    def inference(
+        self,
+        *,
+        token: torch.Tensor,
+        prefix_target: torch.Tensor | None = None,
+        spkr_embed: torch.Tensor | None = None,
+        s_steps: int | None = 10,
+        cfg_alpha: float | None = 2.0,
+        rescale_logits: bool = False,
+        **kwargs,
+    ):
+        b, device = token.size(0), token.device
+        assert b == 1, "Only support batch_size == 1 when inference"
+
+        token = self.token_proj(token)
+        ge = self.cond_mask_spkr_embed(token, spkr_embed)
+        tgt_lens = token.size(1)
+        latent_cond = torch.zeros(b, tgt_lens, self.target_dim).to(device)
+        if prefix_target is not None:
+            latent_cond[:, : prefix_target.size(1), :] = prefix_target
+        sample, trajectory = self.sample(
+            tokens=token, audio=latent_cond, steps=s_steps, alpha=cfg_alpha, g_cond=ge, rescale_logits=rescale_logits
+        )
+        return sample
+
+    def sample(self, tokens, audio, steps, alpha=None, g_cond=None, rescale_logits=False):
+        noise = torch.randn([audio.size(0), audio.shape[1], self.target_dim], device=audio.device)
+        times = torch.linspace(0, 1, steps, device=audio.device)
+
+        def solver(t, z):
+            if alpha is None:
+                output = torch.cat([tokens, audio, z], dim=-1)
+                return self.vectorfield_forward(inputs=output, times=t.unsqueeze(0), self_attn_mask=None, g_cond=g_cond)
+            tokens_empty = torch.zeros(*audio.shape[:2], self.model_dim, device=tokens.device, dtype=tokens.dtype)
+            audio_empty = audio
+            tokens_t = torch.cat([tokens_empty, tokens], dim=0)
+            audio_t = torch.cat([audio_empty, audio], dim=0)
+            audio_noizy_t = torch.cat([z, z], dim=0)
+            t_t = torch.stack([t, t], dim=0)
+            c = g_cond
+            if g_cond is not None:
+                c = torch.cat([g_cond, g_cond], dim=0)
+            output = torch.cat([tokens_t, audio_t, audio_noizy_t], dim=-1)
+            predicted_mix = self.vectorfield_forward(inputs=output, times=t_t, self_attn_mask=None, g_cond=c)
+            predicted_conditioned = predicted_mix[1].unsqueeze(0)
+            predicted_unconditioned = predicted_mix[0].unsqueeze(0)
+
+            prediction = predicted_conditioned + (predicted_conditioned - predicted_unconditioned) * alpha
+            prediction_rescaled = predicted_conditioned.std() * (prediction / prediction.std())
+            if rescale_logits:
+                return prediction_rescaled
+            else:
+                return prediction
+
+        try:
+            from torchdiffeq import odeint
+        except ImportError:
+            raise ImportError("Covo-Audio code2wav requires `torchdiffeq`. Install it with: pip install torchdiffeq")
+        trajectory = odeint(solver, noise, times, atol=1e-5, rtol=1e-5, method="midpoint")
+        return trajectory[-1], trajectory
+
+
+class Token2latentFlowMatchingWithEmbed(Token2latentFlowMatching):
+    def __init__(self, config, **kwargs):
+        super().__init__(config)
+        self.vocab_size = config.token_vocab_size
+        self.token_embedding = EmbeddingTable(
+            num_embeddings=self.vocab_size, embedding_dim=self.token_input_dim, pad_id=self.token_pad_id
+        )
+
+    def inference(
+        self, *, token, prefix_target=None, spkr_embed=None, s_steps=10, cfg_alpha=2, rescale_logits=False, **kwargs
+    ):
+        token = self.token_embedding(token)
+        return super().inference(
+            token=token,
+            prefix_target=prefix_target,
+            spkr_embed=spkr_embed,
+            s_steps=s_steps,
+            cfg_alpha=cfg_alpha,
+            rescale_logits=rescale_logits,
+            **kwargs,
+        )
+
+
+# --------------- Vocoder: Latent -> Waveform (BigVGAN) ---------------
+
+
+class AMPBlock1(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5), causal=True):
+        super().__init__()
+        self.h = h
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], causal=causal)),
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], causal=causal)),
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], causal=causal)),
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)),
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)),
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)),
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+        self.num_layers = len(self.convs1) + len(self.convs2)
+
+        self.activations = nn.ModuleList(
+            [
+                Activation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale))
+                for _ in range(self.num_layers)
+            ]
+        )
+
+    def forward(self, x):
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+            xt = a1(x)
+            xt = c1(xt)
+            xt = a2(xt)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+
+class BigVGANFlowVAE(nn.Module):
+    def __init__(self, h):
+        super().__init__()
+        self.config = h
+        self.h = h
+        causal = h.causal
+        self.hop_size = np.prod(h.downsample_rates)
+
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+
+        self.conv_pre = weight_norm(Conv1d(h.latent_dim, h.upsample_initial_channel, 7, 1, causal=False))
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        weight_norm(
+                            ConvTranspose1d(
+                                h.upsample_initial_channel // (2**i),
+                                h.upsample_initial_channel // (2 ** (i + 1)),
+                                k,
+                                u,
+                                causal=causal,
+                            )
+                        )
+                    ]
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = h.upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
+                self.resblocks.append(AMPBlock1(h, ch, k, d, causal=causal))
+
+        activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale)
+        self.activation_post = Activation1d(activation=activation_post)
+
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, causal=causal))
+
+        for i in range(len(self.ups)):
+            self.ups[i].apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+    def inference_from_latents(self, x, do_sample=True, noise_scale=1.0):
+        if self.h.use_vae and do_sample:
+            assert x.size(1) == self.h.latent_dim * 2, f"Input must be like [B, D, H], got {x.shape}"
+            m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1)
+            x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) * noise_scale
+        else:
+            assert x.size(1) == self.h.latent_dim, f"Input must be like [B, D, H], got {x.shape}"
+
+        x = self.conv_pre(x)
+
+        for i in range(self.num_upsamples):
+            for i_up in range(len(self.ups[i])):
+                x = self.ups[i][i_up](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+
+        x = self.activation_post(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+        return x
+
+
+# --------------- Token2WavDecoder ---------------
+
+
+class Token2WavDecoder(nn.Module):
+    def __init__(self, config, **kwargs):
+        super().__init__()
+        self.config = config
+
+        self.wavegan = BigVGANFlowVAE(config.wavegan)
+        self.wavegan_hop_size = np.prod(config.wavegan.downsample_rates)
+        self.global_mean_var = getattr(config, "global_mean_var", None)
+        if self.global_mean_var is not None:
+            mean_var_data = torch.from_numpy(np.load(self.global_mean_var)).float().squeeze()
+            global_mean, global_var = mean_var_data.chunk(2, 0)
+            self.register_buffer("global_mean", global_mean)
+            self.register_buffer("global_var", global_var)
+
+        self.token2latent = Token2latentFlowMatchingWithEmbed(config.token2latent)
+
+        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)
+
+        token = data["target_token"]
+        prefix_target = None
+        res_prefix_len = 0
+        prompt_token = data.get("prompt_token", None)
+        prompt_latent = data.get("prompt_latent", None)
+        if exists(prompt_token):
+            token = torch.cat([prompt_token, token], dim=1)
+        if exists(prompt_latent):
+            prefix_target = prompt_latent
+            res_prefix_len = prompt_latent.shape[1]
+
+        spkr_embed = None
+        if not zero_spkr:
+            spkr_embed = data.get("spkr_embed", None)
+        res = {"token": token, "prefix_target": prefix_target, "spkr_embed": spkr_embed}
+        return res, res_prefix_len
+
+    @eval_decorator
+    @torch.no_grad()
+    def inference(self, data, **kwargs):
+        infer_data, prefix_len = self.preprocess_infer_data(data)
+        prefix_target = infer_data["prefix_target"]
+        res_latents = self.token2latent.inference(**infer_data, **kwargs)
+        if self.global_mean_var is not None:
+            res_latents = self.global_mean + res_latents * torch.sqrt(self.global_var)
+            if exists(prefix_target):
+                prefix_target = self.global_mean + prefix_target * torch.sqrt(self.global_var)
+        if exists(prefix_target):
+            res_latents[:, : prefix_target.shape[1]] = prefix_target
+        res_latents = res_latents[:, prefix_len:]
+        audio = self.wavegan.inference_from_latents(res_latents.transpose(1, 2), do_sample=False)
+        return audio

vllm_omni/model_executor/models/registry.py

@@ -174,6 +174,32 @@ _OMNI_MODELS = {
         "VoxtralTTSAudioGenerationForConditionalGeneration",
     ),
     "VoxtralTTSAudioTokenizer": ("voxtral_tts", "voxtral_tts_audio_tokenizer", "VoxtralTTSAudioTokenizer"),
+    ## covo_audio
+    "CovoAudioForCausalLM": (
+        "covo_audio",
+        "covo_audio",
+        "CovoAudioForConditionalGeneration",
+    ),
+    "CovoAudioForConditionalGeneration": (
+        "covo_audio",
+        "covo_audio",
+        "CovoAudioForConditionalGeneration",
+    ),
+    "CovoAudioModel": (
+        "covo_audio",
+        "covo_audio",
+        "CovoAudioForConditionalGeneration",
+    ),
+    "CovoAudioLLMModel": (
+        "covo_audio",
+        "covo_audio_llm",
+        "CovoAudioLLMForConditionalGeneration",
+    ),
+    "CovoAudioCode2WavModel": (
+        "covo_audio",
+        "covo_audio_code2wav",
+        "CovoAudioCode2WavForConditionalGeneration",
+    ),
     ## MOSS-TTS-Nano
     "MossTTSNanoForCausalLM": (
         "moss_tts_nano",