fix

Author: hiworldwzj

lightllm/server/httpserver/manager.py

@@ -213,21 +213,6 @@ def tokens(self, prompt, multimodal_params, samping_params: SamplingParams, kwar
             audio_tokens += self.tokenizer.get_audio_token_length(audio)
         return len(prompt_ids) + image_tokens + img_count + audio_tokens + audio_count
 
-    def _calculate_multimodal_tokens(self, multimodal_params: MultimodalParams) -> Tuple[int, int]:
-        image_tokens = 0
-        audio_tokens = 0
-
-        if self.enable_multimodal and self.pd_mode.is_P_or_NORMAL():
-            for img in multimodal_params.images:
-                image_tokens += self.tokenizer.get_image_token_length(img)
-            for audio in multimodal_params.audios:
-                audio_tokens += self.tokenizer.get_audio_token_length(audio)
-        else:
-            image_tokens = len(multimodal_params.images)
-            audio_tokens = len(multimodal_params.audios)
-
-        return image_tokens, audio_tokens
-
     async def loop_for_request(self):
         assert self.args.node_rank > 0
         while True:
@@ -327,15 +312,6 @@ async def generate(
 
             req_status = ReqStatus(group_request_id, multimodal_params, req_objs, start_time)
 
-            # 计算输入 token 使用量统计
-            text_tokens = len(prompt_ids)
-            image_tokens, audio_tokens = self._calculate_multimodal_tokens(multimodal_params)
-            input_usage = {
-                "input_text_tokens": text_tokens,
-                "input_audio_tokens": audio_tokens,
-                "input_image_tokens": image_tokens,
-            }
-
             self.req_id_to_out_inf[group_request_id] = req_status
 
             await self.transfer_to_next_module_or_node(
@@ -350,8 +326,18 @@ async def generate(
                 req_status,
                 request,
             )
+
+            # 计算输入 token 使用量统计
+            image_tokens, audio_tokens = self._count_multimodal_tokens(multimodal_params)
+            text_tokens = len(prompt_ids) - (image_tokens + audio_tokens)
+            input_usage = {
+                "input_text_tokens": text_tokens,
+                "input_audio_tokens": audio_tokens,
+                "input_image_tokens": image_tokens,
+            }
+
             async for sub_req_id, request_output, metadata, finish_status in results_generator:
-                metadata["usage"] = {**input_usage, **metadata.get("usage", {})}
+                metadata["input_usage"] = input_usage
                 yield sub_req_id, request_output, metadata, finish_status
 
         except Exception as e:
@@ -366,6 +352,20 @@ async def generate(
             raise e
         return
 
+    def _count_multimodal_tokens(self, multimodal_params: MultimodalParams) -> Tuple[int, int]:
+        image_tokens = 0
+        audio_tokens = 0
+
+        if self.enable_multimodal and self.pd_mode.is_P_or_NORMAL():
+            for img in multimodal_params.images:
+                if img.token_num is not None:
+                    image_tokens += img.token_num
+            for audio in multimodal_params.audios:
+                if audio.token_num is not None:
+                    audio_tokens += audio.token_num
+
+        return image_tokens, audio_tokens
+
     async def _log_req_header(self, request_headers, group_request_id: int):
 
         x_request_id = request_headers.get("X-Request-Id", "")
@@ -539,8 +539,6 @@ async def _wait_to_token_package(
                     if self.pd_mode == NodeRole.P and is_first_token:
                         metadata["prompt_ids"] = prompt_ids
 
-                    metadata["usage"] = {"output_tokens": out_token_counter}
-
                     prompt_cache_len = metadata.pop("prompt_cache_len", 0)
                     if is_first_token:
                         first_token_cost_ms = (time.time() - start_time) * 1000