Merge branch 'main' into running-eval-suite-skill

Author: yxs

.claude/skills/tune-ci-thresholds/models/s2-pro-v1/config.yaml

@@ -10,14 +10,8 @@
 #   live in that same per-variant dir.
 # - Concurrency: CI uses --concurrency 8 (the only fully tuned conc); the
 #   _VC_*_P95 dicts include {1,2,4,8,16}. Discover reads conc=8 row.
-# - Default venv path: /github/home/omni-s2pro to mirror the real CI venv
-#   (test-s2pro-ci-v1.yaml). Critically, omni-s2pro must NOT have
-#   `openai-whisper` / `whisper-normalizer` installed — CI's omni-s2pro
-#   doesn't, which makes `_get_en_normalizer()` fall back to the
-#   punctuation-strip path. omni-qwen3 (used by qwen3-omni-v1 stages) DOES
-#   have openai-whisper installed and uses the real EnglishTextNormalizer,
-#   so the two venvs produce different WER numbers on identical audio.
-#   Keep them separate; do not unify.
+# - Venv: this host shares the omni-qwen3 venv across models; see
+#   default_venv_python below.
 name: s2-pro-v1
 description: "FishAudio S2-Pro voice-clone TTS (v1 pipeline)"
 hf_model_id: "fishaudio/s2-pro"

benchmarks/tasks/tts.py

@@ -19,13 +19,11 @@
 import time
 import wave
 from dataclasses import dataclass
-from pathlib import Path
 from typing import Protocol
 
 import aiohttp
 import soundfile as sf
 import torch
-import transformers
 from jiwer import process_words
 from tqdm import tqdm
 
@@ -80,57 +78,17 @@ class SampleOutput:
 
 @functools.lru_cache(maxsize=1)
 def _get_en_normalizer():
-    """Lazy-load the English text normalizer.
-
-    Tries whisper_normalizer (standalone pip package) first, then openai-whisper,
-    then the transformers built-in normalizer.
-
-    note (Chenyang): The three fallbacks exist because our deployments don't always
-    have whisper_normalizer installed, whisper's own normalizer lives under a
-    different path depending on the release, and on minimal CI images we rely on
-    the transformers copy bundled with the library.  Keeping all three paths lets
-    the WER numbers stay stable across environments (the official seed-tts-eval
-    reference uses whisper_normalizer, so we prefer it when available).
-    """
-    try:
-        from whisper_normalizer.english import EnglishTextNormalizer
-
-        normalizer = EnglishTextNormalizer()
-        logger.info("Using whisper_normalizer.english.EnglishTextNormalizer")
-        return normalizer
-    except ImportError:
-        logger.debug("whisper_normalizer.english.EnglishTextNormalizer failed")
-
+    """Lazy-load the required English WER normalizer from openai-whisper."""
     try:
         from whisper.normalizers import EnglishTextNormalizer
+    except ImportError as exc:
+        raise RuntimeError(
+            "English WER requires openai-whisper "
+            "(whisper.normalizers.EnglishTextNormalizer). "
+            "Install pinned deps with `uv pip install -e .`."
+        ) from exc
 
-        normalizer = EnglishTextNormalizer()
-        logger.info("Using whisper.normalizers.EnglishTextNormalizer")
-        return normalizer
-    except ImportError:
-        logger.debug("whisper.normalizers.EnglishTextNormalizer failed")
-
-    try:
-        from transformers.models.whisper.english_normalizer import EnglishTextNormalizer
-
-        json_path = (
-            Path(transformers.__file__).parent / "models" / "whisper" / "english.json"
-        )
-        with open(json_path) as f:
-            english_spelling_mapping = json.load(f)
-
-        normalizer = EnglishTextNormalizer(english_spelling_mapping)
-        logger.info(
-            "Using transformers.models.whisper.english_normalizer.EnglishTextNormalizer"
-        )
-        return normalizer
-    except (ImportError, FileNotFoundError) as exc:
-        logger.debug(f"transformers EnglishTextNormalizer failed: {exc}")
-
-    logger.warning(
-        "EnglishTextNormalizer not found; falling back to punctuation-strip normalizer."
-    )
-    return None
+    return EnglishTextNormalizer()
 
 
 def normalize_text(text: str, lang: str) -> str:
@@ -147,15 +105,7 @@ def normalize_text(text: str, lang: str) -> str:
         return text
 
     normalizer = _get_en_normalizer()
-    if normalizer is not None:
-        return normalizer(text)
-
-    for ch in string.punctuation:
-        if ch == "'":
-            continue
-        text = text.replace(ch, "")
-    text = text.replace("  ", " ").strip().lower()
-    return text
+    return normalizer(text)
 
 
 def load_asr_model(lang: str, device: str, generation_mode: str | None = None):

pyproject.toml

@@ -46,7 +46,7 @@ dependencies = [
     "pytest-asyncio>=0.21.0",
     "jiwer",
     "scipy>=1.10.0",
-    "openai-whisper",
+    "openai-whisper==20250625",
     # S2-Pro
     "tiktoken",
     "hydra-core",

sglang_omni_v1/models/fishaudio_s2_pro/bootstrap.py

@@ -106,7 +106,7 @@ def bootstrap_text_model_for_decode(
     audio_decoder: torch.nn.Module,
     semantic_begin_id: int,
     semantic_end_id: int,
-    im_end_id: int,
+    im_end_token_id: int,
     max_batch_size: int,
     num_codebooks: int,
     codebook_size: int,
@@ -119,6 +119,6 @@ def bootstrap_text_model_for_decode(
         codebook_size=codebook_size,
         semantic_begin_id=semantic_begin_id,
         semantic_end_id=semantic_end_id,
-        im_end_id=im_end_id,
+        im_end_token_id=im_end_token_id,
         max_batch_size=max_batch_size,
     )

sglang_omni_v1/models/fishaudio_s2_pro/fish_scheduler.py

@@ -236,7 +236,7 @@ def __init__(
         self, tree_cache: Any, im_end_token_id: int, max_new_tokens: int = 2048
     ):
         self.tree_cache = tree_cache
-        self._im_end_id = int(im_end_token_id)
+        self._im_end_token_id = int(im_end_token_id)
         self._max_new_tokens = int(max_new_tokens)
 
     def update_request(
@@ -250,7 +250,12 @@ def update_request(
             return
 
         if output_token_id is not None:
-            req.output_ids.append(int(output_token_id))
+            semantic_token = int(output_token_id)
+            req.output_ids.append(semantic_token)
+            # Skip caching the terminal slow-AR EOS regardless of req.finished()
+            # semantics: it is not an audio timestep and has no KV to preserve.
+            if semantic_token == self._im_end_token_id:
+                return
             if not req.finished() and req.decode_batch_idx == 0:
                 self.tree_cache.cache_unfinished_req(req)
 
@@ -265,7 +270,7 @@ def is_finished(
         if semantic_token is None and data.previous_semantic_tokens:
             semantic_token = int(data.previous_semantic_tokens[-1])
 
-        if semantic_token == self._im_end_id:
+        if semantic_token == self._im_end_token_id:
             return True
 
         max_tok = data.max_new_tokens or self._max_new_tokens
@@ -418,8 +423,20 @@ def emit_finished(self, finished: list[SchedulerRequest]) -> None:
         for request in finished:
             data = request.data
             data.output_ids = list(data.req.output_ids)
-            result = self._result_adapter(data)
             t_submit = self._submit_times.pop(request.request_id, None)
+            if not data.output_codes:
+                self.outbox.put(
+                    OutgoingMessage(
+                        request_id=request.request_id,
+                        type="error",
+                        data=ValueError(
+                            f"Request {request.request_id}: "
+                            "S2-Pro generated no audio codec tokens"
+                        ),
+                    )
+                )
+                continue
+            result = self._result_adapter(data)
             if t_submit is not None and isinstance(result.data, dict):
                 result.data["engine_time_s"] = time.perf_counter() - t_submit
             self.outbox.put(

sglang_omni_v1/models/fishaudio_s2_pro/model_runner.py

@@ -10,13 +10,44 @@
 from sglang_omni_v1.model_runner.base import ModelRunner
 
 
+def collect_s2pro_step_outputs(
+    result: Any,
+    requests: list,
+    *,
+    output_codes: torch.Tensor,
+    output_semantic_ids: torch.Tensor,
+    im_end_token_id: int,
+) -> None:
+    batch_size = len(requests)
+    if batch_size == 0:
+        return
+
+    result.next_token_ids = output_semantic_ids[:batch_size].clone()
+    semantic_tokens = output_semantic_ids[:batch_size].tolist()
+
+    for row_idx, sched_req in enumerate(requests):
+        data = sched_req.data
+        if data.req.is_chunked > 0:
+            continue
+
+        semantic_token = semantic_tokens[row_idx]
+        if semantic_token == im_end_token_id:
+            continue
+
+        codes = output_codes[row_idx].unsqueeze(-1).clone()
+        data.last_codebook_values = codes[1:, 0].clone()
+        data.previous_semantic_tokens.append(semantic_token)
+        data.output_codes.append(codes)
+
+
 class FishS2ProModelRunner(ModelRunner):
     """Fish TTS runner with unified forward-owned decode and persistent buffers."""
 
     def __init__(self, tp_worker: Any, output_processor: Any):
         super().__init__(tp_worker, output_processor)
         self._semantic_begin_id = int(self.model._semantic_begin_id)
         self._semantic_end_id = int(self.model._semantic_end_id)
+        self._im_end_token_id = int(self.model._im_end_token_id)
 
     def prepare_prefill(self, forward_batch, schedule_batch, requests):
         del schedule_batch
@@ -117,19 +148,10 @@ def _build_prefill_input_embeds(
         return text_embeds
 
     def _collect_step_outputs(self, result: Any, requests: list) -> None:
-        batch_size = len(requests)
-        if batch_size == 0:
-            return
-
-        result.next_token_ids = self.model._output_semantic_ids[:batch_size].clone()
-
-        for row_idx, sched_req in enumerate(requests):
-            data = sched_req.data
-            req = data.req
-            if req.is_chunked > 0:
-                continue
-
-            codes = self.model._output_codes[row_idx].unsqueeze(-1).clone()
-            data.last_codebook_values = codes[1:, 0].clone()
-            data.previous_semantic_tokens.append(int(codes[0, -1].item()))
-            data.output_codes.append(codes)
+        collect_s2pro_step_outputs(
+            result,
+            requests,
+            output_codes=self.model._output_codes,
+            output_semantic_ids=self.model._output_semantic_ids,
+            im_end_token_id=self._im_end_token_id,
+        )

sglang_omni_v1/models/fishaudio_s2_pro/request_builders.py

@@ -101,11 +101,12 @@ def build_sglang_tts_request(
 
 
 def apply_tts_result(state: S2ProState, result: S2ProSGLangRequestData) -> None:
-    if result.output_codes:
-        state.output_codes = torch.cat(result.output_codes, dim=1)
-        state.completion_tokens = state.output_codes.shape[1]
-    else:
-        state.output_codes = None
+    assert result.output_codes, (
+        "apply_tts_result expects non-empty output_codes; "
+        "FishScheduler.emit_finished must filter immediate-EOS cases"
+    )
+    state.output_codes = torch.cat(result.output_codes, dim=1)
+    state.completion_tokens = state.output_codes.shape[1]
     state.prompt_tokens = len(result.input_ids) if result.input_ids is not None else 0
 
 

sglang_omni_v1/models/fishaudio_s2_pro/sglang_model.py

@@ -242,7 +242,7 @@ def setup_vq_decode(
         codebook_size: int,
         semantic_begin_id: int,
         semantic_end_id: int,
-        im_end_id: int,
+        im_end_token_id: int,
         max_batch_size: int,
     ) -> None:
         """Attach audio decoder and allocate persistent GPU buffers."""
@@ -254,6 +254,7 @@ def setup_vq_decode(
         self._num_codebooks = num_codebooks
         self._semantic_begin_id = semantic_begin_id
         self._semantic_end_id = semantic_end_id
+        self._im_end_token_id = int(im_end_token_id)
 
         # Shared codebook embedding from audio decoder (for VQ input combination)
         self._vq_codebook_embeddings = audio_decoder.codebook_embeddings
@@ -271,7 +272,7 @@ def setup_vq_decode(
             (self.vocab_size,), -float("inf"), device=device, dtype=torch.bfloat16
         )
         bias[semantic_begin_id : semantic_end_id + 1] = 0.0
-        bias[im_end_id] = 0.0
+        bias[im_end_token_id] = 0.0
         self._semantic_bias = bias
 
         # Output buffers: written by _decode_codebooks, read by ModelRunner

sglang_omni_v1/models/fishaudio_s2_pro/stages.py

@@ -248,7 +248,7 @@ def create_sglang_tts_engine_executor(
         audio_decoder=audio_decoder,
         semantic_begin_id=adapter.semantic_begin_id,
         semantic_end_id=adapter.semantic_end_id,
-        im_end_id=adapter.eos_token_ids[0],
+        im_end_token_id=adapter.eos_token_ids[0],
         max_batch_size=server_args.max_running_requests,
         num_codebooks=num_codebooks,
         codebook_size=codebook_size,

sglang_omni_v1/models/qwen3_omni/request_builders.py

@@ -18,6 +18,7 @@
 from sglang_omni_v1.proto import StagePayload
 from sglang_omni_v1.scheduling.messages import OutgoingMessage
 from sglang_omni_v1.scheduling.sglang_backend import SGLangARRequestData
+from sglang_omni_v1.scheduling.types import ARRequestData
 
 IMAGE_STAGE = "image_encoder"
 AUDIO_STAGE = "audio_encoder"
@@ -32,10 +33,6 @@ class EncoderRequestData:
     skip_result: dict[str, Any] | None = None
 
 
-class ARRequestData:
-    """AR request data — base for SGLangARRequestData."""
-
-
 def build_encoder_request(
     state: PipelineState, *, stage_name: str
 ) -> EncoderRequestData: