[Feat]Support Nonverbal Tags in OmniVoice (#3968)

Signed-off-by: sphinxkkkbc <binchengkang8@gmail.com>
Co-authored-by: sphinxkkkbc <binchengkang8@gmail.com>

tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import pytest
+from huggingface_hub import hf_hub_download
+from tokenizers import Tokenizer as HFTokenizer
+
+from vllm_omni.diffusion.models.omnivoice.pipeline_omnivoice import (
+    _tokenize_with_nonverbal_tags,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _ids(tok: HFTokenizer, s: str) -> list[int]:
+    return tok.encode(s).ids
+
+
+class TestNonVerbalTags:
+    @classmethod
+    def setup_class(cls):
+        tokenizer_path = hf_hub_download(repo_id="k2-fsa/OmniVoice", filename="tokenizer.json")
+        cls.tokenizer = HFTokenizer.from_file(tokenizer_path)
+
+    def test_plain_text_fallback(self):
+        tok = self.tokenizer
+        text = "hello world"
+        assert _tokenize_with_nonverbal_tags(text, tok) == _ids(tok, text)
+
+    def test_single_nonverbal_tag_at_beginning(self):
+        tok = self.tokenizer
+        text = "[laughter]HelloWorld"
+        expected = _ids(tok, "[laughter]") + _ids(tok, "Hello") + _ids(tok, "World")
+        assert _tokenize_with_nonverbal_tags(text, tok) == expected
+
+    def test_single_nonverbal_tag_in_middle(self):
+        tok = self.tokenizer
+        text = "Hello[laughter]world"
+        expected = _ids(tok, "Hello") + _ids(tok, "[laughter]") + _ids(tok, "world")
+        assert _tokenize_with_nonverbal_tags(text, tok) == expected
+
+    def test_multiple_nonverbal_tags(self):
+        tok = self.tokenizer
+        text = "[sigh]a[question-en]b[surprise-oh]"
+        expected = (
+            _ids(tok, "[sigh]")
+            + _ids(tok, "a")
+            + _ids(tok, "[question-en]")
+            + _ids(tok, "b")
+            + _ids(tok, "[surprise-oh]")
+        )
+        assert _tokenize_with_nonverbal_tags(text, tok) == expected
+
+    def test_tag_only(self):
+        tok = self.tokenizer
+        text = "[question-en]"
+        assert _tokenize_with_nonverbal_tags(text, tok) == _ids(tok, text)
+
+    def test_unknown_tag_not_special_cased(self):
+        tok = self.tokenizer
+        text = "x[unknown-tag]y"
+        assert _tokenize_with_nonverbal_tags(text, tok) == _ids(tok, text)

vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py

@@ -13,6 +13,7 @@ from __future__ import annotations
 
 import json
 import os
+import re
 from collections.abc import Iterable
 from typing import ClassVar
 
@@ -53,6 +54,76 @@ def get_omnivoice_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
+def _combine_text(text, ref_text: str | None = None) -> str:
+    # combine with reference text if not None
+    if ref_text:
+        full_text = ref_text.strip() + " " + text.strip()
+    else:
+        full_text = text.strip()
+
+    # filter out newline / carriage-return characters
+    full_text = re.sub(r"[\r\n]+", "", full_text)
+
+    # replace Chinese parentheses with English ones
+    full_text = full_text.replace("\uff08", "(").replace("\uff09", ")")
+
+    # collapse consecutive spaces / tabs into a single space
+    full_text = re.sub(r"[ \t]+", " ", full_text)
+
+    # remove spaces around chinese characters
+    chinese_range = r"[\u4e00-\u9fff]"
+    pattern = rf"(?<={chinese_range})\s+|\s+(?={chinese_range})"
+    full_text = re.sub(pattern, "", full_text)
+
+    return full_text
+
+
+_NONVERBAL_PATTERN = re.compile(
+    r"\[(laughter|sigh|confirmation-en|question-en|question-ah|question-oh|"
+    r"question-ei|question-yi|surprise-ah|surprise-oh|surprise-wa|"
+    r"surprise-yo|dissatisfaction-hnn)\]"
+)
+
+
+def _tokenize_with_nonverbal_tags(text: str, tokenizer) -> list[int]:
+    """Tokenize text containing non-verbal tags, handling each tag independently.
+
+    Non-verbal tags are tokenized standalone to guarantee consistent token
+    IDs regardless of surrounding language context (Chinese, English, etc.).
+
+    Args:
+        text: Full text string potentially containing non-verbal tags.
+        tokenizer: HuggingFace text tokenizer instance.
+    Returns:
+        Token IDs list of length seq_len.
+    """
+    parts = []
+    last_end = 0
+    for m in _NONVERBAL_PATTERN.finditer(text):
+        if m.start() > last_end:
+            segment = text[last_end : m.start()]
+            ids = tokenizer.encode(segment)
+            if ids:
+                parts.append(ids)
+        tag_ids = tokenizer.encode(m.group())
+        if tag_ids:
+            parts.append(tag_ids)
+        last_end = m.end()
+    if last_end < len(text):
+        segment = text[last_end:]
+        ids = tokenizer.encode(segment)
+        if ids:
+            parts.append(ids)
+
+    if not parts:
+        return tokenizer.encode(text).ids
+    else:
+        combined = []
+        for p in parts:
+            combined.extend(p.ids)
+    return combined
+
+
 class OmniVoicePipeline(nn.Module, SupportAudioOutput):
     """OmniVoice text-to-speech pipeline for the diffusion engine.
 
@@ -212,14 +283,13 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput):
         target_len = max(1, int(target_len))
 
         # Build text prompt with control tokens
-        style = f"<|denoise|><|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>"
-        if ref_text:
-            full_text = f"{ref_text} {text}"
-        else:
-            full_text = text
-        full_prompt = f"{style}<|text_start|>{full_text}<|text_end|>"
-        encoding = self.tokenizer.encode(full_prompt)
-        text_tokens = torch.tensor(encoding.ids, dtype=torch.long, device=device)
+        style_text = f"<|denoise|><|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>"
+        full_text = _combine_text(ref_text=ref_text, text=text)
+        wrapped_text = f"<|text_start|>{full_text}<|text_end|>"
+        style_tokens = self.tokenizer.encode(style_text).ids
+        text_tokens = _tokenize_with_nonverbal_tags(wrapped_text, self.tokenizer)
+        encoding_ids = style_tokens + text_tokens
+        text_tokens = torch.tensor(encoding_ids, dtype=torch.long, device=device)
         text_len = text_tokens.shape[0]
 
         # Encode reference audio tokens if provided (with voice caching)