[Not for merge] Speculative Decoding Cosyvoice LLM example

cosyvoice/cli/model.py

@@ -422,6 +422,49 @@ def __init__(self,
         # FSQ silent and breath token
         self.silent_tokens = [1, 2, 28, 29, 55, 248, 494, 2241, 2242, 2322, 2323]
 
+    def load_vllm_spec(self, model_dir, draft_model_dir, gpu_memory_utilization=0.6):
+        """Load CosyVoice3 + DSpark speculative decoding (vllm V1 API, token-based).
+
+        `model_dir` is the CosyVoice3 pretrained model directory (same as used to load
+        the model, e.g. 'pretrained_models/Fun-CosyVoice3-0.5B').  The HuggingFace export
+        is created automatically under <model_dir>/hf_spec the first time (like
+        export_cosyvoice2_vllm() for load_vllm()).
+        `draft_model_dir` is a speculators-format DSpark checkpoint (local path or HF repo).
+
+        After loading, self.llm.vllm_spec is used automatically inside tts() calls.
+        """
+        import json
+        from transformers import AutoTokenizer
+        from vllm import LLM
+        from cosyvoice.utils.file_utils import export_cosyvoice3_vllm_spec
+
+        hf_llm_dir = os.path.join(model_dir, 'CosyVoice-BlankEN')
+        hf_model_dir = os.path.join(model_dir, 'hf_spec')
+        export_cosyvoice3_vllm_spec(self.llm, hf_model_dir, hf_llm_dir, self.device)
+
+        with open(os.path.join(hf_model_dir, 'cosyvoice3_metadata.json')) as f:
+            metadata = json.load(f)
+        with open(os.path.join(draft_model_dir, 'config.json')) as f:
+            draft_cfg = json.load(f)
+        method = draft_cfg.get('speculators_model_type', 'dspark')
+        num_spec_tokens = draft_cfg.get('block_size', 8) - 1
+
+        self.llm.vllm_spec_tokenizer = AutoTokenizer.from_pretrained(hf_model_dir)
+        self.llm.vllm_spec_speech_offset = metadata['speech_token_offset']
+        self.llm.vllm_spec = LLM(
+            model=hf_model_dir,
+            speculative_config={
+                'model': draft_model_dir,
+                'method': method,
+                'num_speculative_tokens': num_spec_tokens,
+                'draft_sample_method': 'probabilistic',
+                'draft_apply_repetition_penalty': True,
+            },
+            gpu_memory_utilization=gpu_memory_utilization,
+        )
+        # free PyTorch model weights; vllm manages GPU memory from here
+        del self.llm.llm.model.model.layers
+
     def token2wav(self, token, prompt_token, prompt_feat, embedding, token_offset, uuid, stream=False, finalize=False, speed=1.0):
         with torch.cuda.amp.autocast(self.fp16):
             tts_mel, _ = self.flow.inference(token=token.to(self.device, dtype=torch.int32),

cosyvoice/llm/llm.py

@@ -498,6 +498,10 @@ def inference(
         max_len = int((text_len - prompt_text_len) * max_token_text_ratio)
 
         # 5. step by step decode
+        # speculative vllm path: token-based V1 API (CosyVoice3 only)
+        if hasattr(self, 'vllm_spec'):
+            yield from self._inference_spec(text, prompt_speech_token, min_len, max_len)
+            return
         for token in self.inference_wrapper(lm_input, sampling, min_len, max_len, uuid):
             yield token
 
@@ -703,4 +707,37 @@ def __init__(
         self.stop_token_ids = [speech_token_size + i for i in range(200)]
         self.vllm_output_queue = {}
         if online_feature is True:
-            self.speech_token_extractor = SpeechTokenExtractor(model_path=os.path.join(onnx_path, 'speech_tokenizer_v3.batch.onnx'))
+            self.speech_token_extractor = SpeechTokenExtractor(model_path=os.path.join(onnx_path, 'speech_tokenizer_v3.batch.onnx'))
+
+    def _inference_spec(self, text: torch.Tensor, prompt_speech_token: torch.Tensor,
+                        min_len: int, max_len: int):
+        """Token-based speculative decoding via vllm V1 LLM (set by load_vllm_spec).
+
+        `text` is the concatenated [prompt_text, target_text] as Qwen2 token IDs.
+        `prompt_speech_token` holds voice-clone speech token IDs (0..speech_token_size).
+        Output: yields integer speech token IDs.
+        """
+        from vllm import SamplingParams
+        offset = self.vllm_spec_speech_offset
+        # only pass the primary eos stop token (vllm caps stop_token_ids at 128);
+        # the yield loop below also filters any token outside [0, speech_token_size)
+        stop_ids = [offset + self.eos_token]
+
+        text_str = self.vllm_spec_tokenizer.decode(text[0].tolist(), skip_special_tokens=False)
+        speech_str = ''.join(f'<|s_{t}|>' for t in prompt_speech_token[0].tolist()) if prompt_speech_token.shape[1] > 0 else ''
+
+        chat = [{'role': 'user', 'content': text_str}]
+        if speech_str:
+            chat.append({'role': 'assistant', 'content': speech_str})
+        prompt = self.vllm_spec_tokenizer.apply_chat_template(
+            chat, tokenize=False, continue_final_message=True)
+
+        params = SamplingParams(
+            temperature=0.8, top_p=0.95, top_k=15, repetition_penalty=1.1,
+            min_tokens=min_len, max_tokens=max_len, stop_token_ids=stop_ids)
+        output = self.vllm_spec.generate([prompt], params, use_tqdm=False)[0]
+        for token_id in output.outputs[0].token_ids:
+            speech_id = token_id - offset
+            if speech_id < 0 or speech_id >= self.speech_token_size:
+                break
+            yield speech_id

cosyvoice/utils/file_utils.py

@@ -89,6 +89,141 @@ def convert_onnx_to_trt(trt_model, trt_kwargs, onnx_model, fp16):
 
 
 # NOTE do not support bistream inference as only speech token embedding/head is kept
+_COSYVOICE3_TEXT_SPECIAL_TOKENS = [
+    '<|im_start|>', '<|im_end|>', '<|endofprompt|>',
+    '[breath]', '<strong>', '</strong>', '[noise]', '[laughter]', '[cough]',
+    '[clucking]', '[accent]', '[quick_breath]', '<laughter>', '</laughter>',
+    '[hissing]', '[sigh]', '[vocalized-noise]', '[lipsmack]', '[mn]', '<|endofsystem|>',
+    '[AA]', '[AA0]', '[AA1]', '[AA2]', '[AE]', '[AE0]', '[AE1]', '[AE2]',
+    '[AH]', '[AH0]', '[AH1]', '[AH2]', '[AO]', '[AO0]', '[AO1]', '[AO2]',
+    '[AW]', '[AW0]', '[AW1]', '[AW2]', '[AY]', '[AY0]', '[AY1]', '[AY2]',
+    '[B]', '[CH]', '[D]', '[DH]', '[EH]', '[EH0]', '[EH1]', '[EH2]',
+    '[ER]', '[ER0]', '[ER1]', '[ER2]', '[EY]', '[EY0]', '[EY1]', '[EY2]',
+    '[F]', '[G]', '[HH]', '[IH]', '[IH0]', '[IH1]', '[IH2]',
+    '[IY]', '[IY0]', '[IY1]', '[IY2]', '[JH]', '[K]', '[L]', '[M]', '[N]',
+    '[NG]', '[OW]', '[OW0]', '[OW1]', '[OW2]', '[OY]', '[OY0]', '[OY1]', '[OY2]',
+    '[P]', '[R]', '[S]', '[SH]', '[T]', '[TH]', '[UH]', '[UH0]', '[UH1]', '[UH2]',
+    '[UW]', '[UW0]', '[UW1]', '[UW2]', '[V]', '[W]', '[Y]', '[Z]', '[ZH]',
+    '[a]', '[ai]', '[an]', '[ang]', '[ao]', '[b]', '[c]', '[ch]', '[d]', '[e]',
+    '[ei]', '[en]', '[eng]', '[f]', '[g]', '[h]', '[i]', '[ian]', '[in]', '[ing]',
+    '[iu]', '[ià]', '[iàn]', '[iàng]', '[iào]', '[iá]', '[ián]', '[iáng]', '[iáo]',
+    '[iè]', '[ié]', '[iòng]', '[ióng]', '[iù]', '[iú]', '[iā]', '[iān]', '[iāng]',
+    '[iāo]', '[iē]', '[iě]', '[iōng]', '[iū]', '[iǎ]', '[iǎn]', '[iǎng]', '[iǎo]',
+    '[iǒng]', '[iǔ]', '[j]', '[k]', '[l]', '[m]', '[n]', '[o]', '[ong]', '[ou]',
+    '[p]', '[q]', '[r]', '[s]', '[sh]', '[t]', '[u]', '[uang]', '[ue]', '[un]',
+    '[uo]', '[uà]', '[uài]', '[uàn]', '[uàng]', '[uá]', '[uái]', '[uán]', '[uáng]',
+    '[uè]', '[ué]', '[uì]', '[uí]', '[uò]', '[uó]', '[uā]', '[uāi]', '[uān]',
+    '[uāng]', '[uē]', '[uě]', '[uī]', '[uō]', '[uǎ]', '[uǎi]', '[uǎn]', '[uǎng]',
+    '[uǐ]', '[uǒ]', '[vè]', '[w]', '[x]', '[y]', '[z]', '[zh]',
+    '[à]', '[ài]', '[àn]', '[àng]', '[ào]', '[á]', '[ái]', '[án]', '[áng]', '[áo]',
+    '[è]', '[èi]', '[èn]', '[èng]', '[èr]', '[é]', '[éi]', '[én]', '[éng]', '[ér]',
+    '[ì]', '[ìn]', '[ìng]', '[í]', '[ín]', '[íng]', '[ò]', '[òng]', '[òu]',
+    '[ó]', '[óng]', '[óu]', '[ù]', '[ùn]', '[ú]', '[ún]', '[ā]', '[āi]', '[ān]',
+    '[āng]', '[āo]', '[ē]', '[ēi]', '[ēn]', '[ēng]', '[ě]', '[ěi]', '[ěn]',
+    '[ěng]', '[ěr]', '[ī]', '[īn]', '[īng]', '[ō]', '[ōng]', '[ōu]', '[ū]',
+    '[ūn]', '[ǎ]', '[ǎi]', '[ǎn]', '[ǎng]', '[ǎo]', '[ǐ]', '[ǐn]', '[ǐng]',
+    '[ǒ]', '[ǒng]', '[ǒu]', '[ǔ]', '[ǔn]', '[ǘ]', '[ǚ]', '[ǜ]',
+]
+
+# Chat template for CosyVoice3 speculative vllm: <|sos|>{text}<|task_id|>{speech_tokens}
+_COSYVOICE3_CHAT_TEMPLATE = (
+    "{%- for message in messages %}"
+    "{%- if message['role'] == 'user' %}{{- '<|sos|>' + message['content'] + '<|task_id|>' }}"
+    "{%- elif message['role'] == 'assistant' %}{{- message['content']}}"
+    "{%- endif %}{%- endfor %}"
+)
+
+
+def export_cosyvoice3_vllm_spec(model, model_path, hf_llm_dir, device):
+    """Export CosyVoice3LM to HuggingFace format compatible with DSpark speculative decoding.
+
+    Extends the Qwen2 tokenizer with CosyVoice3 text special tokens and speech tokens,
+    then splices speech_embedding / llm_decoder into the model's embed_tokens / lm_head
+    at offset `text_vocab_size`. This produces the same token-ID layout expected by the
+    DSpark draft model (yuekai/cosyvoice3_llm_dspark): speech token N lives at token ID
+    `text_vocab_size + N`.
+
+    Mirrors export_cosyvoice2_vllm() in style: skips if model_path already exists.
+    """
+    if os.path.exists(model_path):
+        return
+
+    from transformers import AutoTokenizer
+    import math
+
+    dtype = torch.bfloat16
+    tokenizer = AutoTokenizer.from_pretrained(hf_llm_dir, trust_remote_code=True)
+    tokenizer.add_special_tokens({
+        'eos_token': '<|endoftext|>',
+        'pad_token': '<|endoftext|>',
+        'additional_special_tokens': _COSYVOICE3_TEXT_SPECIAL_TOKENS,
+    })
+    text_vocab_size = len(tokenizer)
+
+    speech_token_size = model.speech_embedding.num_embeddings  # includes 200 special
+    base_speech_token_size = model.speech_token_size            # e.g. 2512 or 6561
+
+    speech_tokens = [f'<|s_{i}|>' for i in range(speech_token_size)]
+    speech_tokens[base_speech_token_size + 0] = '<|sos|>'
+    speech_tokens[base_speech_token_size + 1] = '<|eos1|>'
+    speech_tokens[base_speech_token_size + 2] = '<|task_id|>'
+    speech_tokens[base_speech_token_size + 3] = '<|fill|>'
+    tokenizer.add_tokens(speech_tokens)
+    tokenizer.chat_template = _COSYVOICE3_CHAT_TEMPLATE
+
+    # Pad vocab to multiple of 128 for efficiency
+    new_vocab = len(tokenizer)
+    padded_vocab = math.ceil(new_vocab / 128) * 128
+
+    qwen_model = model.llm.model
+    qwen_model.resize_token_embeddings(padded_vocab)
+    qwen_model.to(dtype).to(device)
+
+    with torch.no_grad():
+        # Splice speech_embedding into embed_tokens at [text_vocab_size:]
+        src = min(model.speech_embedding.weight.shape[0], speech_token_size)
+        qwen_model.get_input_embeddings().weight[text_vocab_size:text_vocab_size + src] = \
+            model.speech_embedding.weight[:src].to(dtype)
+
+        # Build new lm_head: text part zeroed (not generated), speech part = llm_decoder
+        has_bias = model.llm_decoder.bias is not None
+        new_lm_head = torch.nn.Linear(
+            qwen_model.config.hidden_size, padded_vocab, bias=has_bias, device=device, dtype=dtype)
+        new_lm_head.weight.data.zero_()
+        if has_bias:
+            new_lm_head.bias.data.fill_(float('-inf'))
+        # Copy text part from original lm_head so text logits stay valid (optional but clean)
+        orig_head = qwen_model.lm_head
+        copy_text = min(orig_head.weight.shape[0], text_vocab_size)
+        new_lm_head.weight[:copy_text] = orig_head.weight[:copy_text].to(dtype)
+        # Copy llm_decoder into speech part
+        dec_size = min(model.llm_decoder.weight.shape[0], speech_token_size)
+        new_lm_head.weight[text_vocab_size:text_vocab_size + dec_size] = \
+            model.llm_decoder.weight[:dec_size].to(dtype)
+        if has_bias:
+            new_lm_head.bias[text_vocab_size:text_vocab_size + dec_size] = \
+                model.llm_decoder.bias[:dec_size].to(dtype)
+        qwen_model.lm_head = new_lm_head
+
+    eos_id = text_vocab_size + base_speech_token_size + 1
+    qwen_model.config.vocab_size = padded_vocab
+    qwen_model.config.tie_word_embeddings = False
+    qwen_model.generation_config.eos_token_id = eos_id
+    qwen_model.generation_config.pad_token_id = eos_id
+
+    os.makedirs(model_path, exist_ok=True)
+    qwen_model.save_pretrained(model_path)
+    tokenizer.save_pretrained(model_path)
+    json.dump({
+        'text_vocab_size': text_vocab_size,
+        'base_speech_token_size': base_speech_token_size,
+        'embedding_size': speech_token_size,
+        'padded_vocab_size': padded_vocab,
+        'eos_token_id': eos_id,
+        'speech_token_offset': text_vocab_size,
+    }, open(os.path.join(model_path, 'cosyvoice3_metadata.json'), 'w'), indent=2)
+
+
 def export_cosyvoice2_vllm(model, model_path, device):
     if os.path.exists(model_path):
         return

vllm_example.py

@@ -30,9 +30,36 @@ def cosyvoice3_example():
             continue
 
 
+def cosyvoice3_spec_example():
+    """CosyVoice3 + DSpark speculative decoding.
+
+    Setup (one-time):
+        pip install vllm-omni==0.25.1
+        git clone -b dspark-draft-sampling-mirrors https://github.com/yuekaizhang/vllm
+        export PYTHONPATH=/path/to/yuekaizhang/vllm:$PYTHONPATH
+
+    Then run:
+        python vllm_example.py
+    """
+    DRAFT_MODEL = 'yuekai/cosyvoice3_llm_dspark'  # or local path to checkpoint_best/
+
+    cosyvoice = AutoModel(model_dir='pretrained_models/Fun-CosyVoice3-0.5B', load_trt=True, fp16=False)
+    # HF export is created automatically under pretrained_models/Fun-CosyVoice3-0.5B/hf_spec
+    cosyvoice.model.load_vllm_spec('pretrained_models/Fun-CosyVoice3-0.5B', DRAFT_MODEL)
+
+    for i in tqdm(range(10)):
+        set_all_random_seed(i)
+        for j, result in enumerate(cosyvoice.inference_zero_shot(
+                '收到好友从远方寄来的生日礼物，那份意外的惊喜与深深的祝福让我心中充满了甜蜜的快乐，笑容如花儿般绽放。',
+                'You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。',
+                './asset/zero_shot_prompt.wav', stream=False)):
+            result['tts_speech'].numpy().tofile(f'spec_{i}_{j}.pcm')
+
+
 def main():
     # cosyvoice2_example()
     cosyvoice3_example()
+    # cosyvoice3_spec_example()
 
 
 if __name__ == '__main__':