Fix S2-Pro v1 streaming usage metrics

Author: Ratish1

sglang_omni_v1/models/fishaudio_s2_pro/fish_scheduler.py

@@ -416,6 +416,11 @@ def emit_finished(self, finished: list[SchedulerRequest]) -> None:
         for request in finished:
             data = request.data
             data.output_ids = list(data.req.output_ids)
+            assert request.finish_time is not None
+            data.stage_payload.data["engine_time_s"] = max(
+                request.finish_time - request.arrival_time,
+                1e-6,
+            )
             result = self._result_adapter(data)
             self.outbox.put(
                 OutgoingMessage(

sglang_omni_v1/models/fishaudio_s2_pro/request_builders.py

@@ -109,6 +109,17 @@ def apply_tts_result(state: S2ProState, result: S2ProSGLangRequestData) -> None:
     state.prompt_tokens = len(result.input_ids) if result.input_ids is not None else 0
 
 
+def build_tts_usage(state: S2ProState) -> dict[str, Any]:
+    usage = {
+        "prompt_tokens": int(state.prompt_tokens),
+        "completion_tokens": int(state.completion_tokens),
+        "total_tokens": int(state.prompt_tokens + state.completion_tokens),
+    }
+    if state.engine_time_s > 0:
+        usage["engine_time_s"] = round(float(state.engine_time_s), 6)
+    return usage
+
+
 def make_tts_scheduler_adapters(*, tokenizer: Any):
     """Build model-specific StagePayload <-> scheduler adapters for Fish TTS."""
 
@@ -126,10 +137,13 @@ def result_adapter(data: S2ProSGLangRequestData) -> StagePayload:
         payload = data.stage_payload
         state = S2ProState.from_dict(payload.data)
         apply_tts_result(state, data)
+        state.engine_time_s = float(payload.data["engine_time_s"])
+        result_data = state.to_dict()
+        result_data["usage"] = build_tts_usage(state)
         return StagePayload(
             request_id=payload.request_id,
             request=payload.request,
-            data=state.to_dict(),
+            data=result_data,
         )
 
     return request_builder, result_adapter

sglang_omni_v1/models/fishaudio_s2_pro/stages.py

@@ -291,9 +291,11 @@ def _store_audio(
         state: S2ProState,
         audio_np: torch.Tensor,
     ) -> StagePayload:
+        usage = payload.data["usage"]
         state.audio_samples = audio_np
         state.sample_rate = codec.sample_rate
         payload = store_state(payload, state)
+        payload.data["usage"] = usage
         payload.data["audio_data"] = audio_np.tolist()
         payload.data["sample_rate"] = codec.sample_rate
         payload.data["modality"] = "audio"

tests/test_v1_fish_vocoder_batch.py

@@ -36,6 +36,18 @@ def _payload(request_id: str, code_len: int) -> StagePayload:
     )
 
 
+def _run_vocoder_request(scheduler, payload: StagePayload) -> StagePayload:
+    thread = threading.Thread(target=scheduler.start, daemon=True)
+    thread.start()
+    try:
+        scheduler.inbox.put(IncomingMessage(payload.request_id, "new_request", payload))
+        output = scheduler.outbox.get(timeout=2.0)
+        return output.data
+    finally:
+        scheduler.stop()
+        thread.join(timeout=2.0)
+
+
 def test_fish_vocoder_uses_simple_scheduler_batch_path(monkeypatch) -> None:
     codec = _FakeCodec()
     monkeypatch.setattr(stages, "_resolve_checkpoint", lambda model_path: model_path)
@@ -69,3 +81,31 @@ def test_fish_vocoder_uses_simple_scheduler_batch_path(monkeypatch) -> None:
     assert len(outputs["req-long"].data["audio_data"]) == 12
     assert outputs["req-short"].data["audio_data"] == [1.0] * 8
     assert outputs["req-long"].data["audio_data"] == [2.0] * 12
+
+
+def test_fish_vocoder_preserves_existing_usage(monkeypatch) -> None:
+    codec = _FakeCodec()
+    monkeypatch.setattr(stages, "_resolve_checkpoint", lambda model_path: model_path)
+    monkeypatch.setattr(stages, "_load_codec", lambda checkpoint, device: codec)
+
+    payload = _payload("req-usage", 2)
+    usage = {
+        "prompt_tokens": 3,
+        "completion_tokens": 2,
+        "total_tokens": 5,
+        "engine_time_s": 0.25,
+    }
+    payload.data["usage"] = usage
+    scheduler = stages.create_vocoder_executor(
+        "unused",
+        device="cpu",
+        max_batch_size=1,
+        max_batch_wait_ms=1,
+    )
+
+    output = _run_vocoder_request(scheduler, payload)
+
+    assert output.data["usage"] == usage
+    assert output.data["audio_data"] == [1.0] * 8
+    assert output.data["sample_rate"] == 44100
+    assert output.data["modality"] == "audio"