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"