[S2-Pro]: Fix S2-Pro terminal EOS audio frame (#377)

Co-authored-by: zhaochenyang20 <zhaochen20@outlook.com>

Author: Ratish1

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,