Update kimi_vl/serve/inference.py

kimi_vl/serve/inference.py

@@ -4,7 +4,7 @@ from threading import Thread
 from typing import List, Optional
 
 import torch
-import spaces
+
 from transformers import (
     AutoModelForCausalLM,
     AutoProcessor,
@@ -73,6 +73,7 @@ def preprocess(
     messages: list[dict],
     processor,
     sft_format: Optional[str] = "kimi-vl",
+    override_system_prompt = "",
 ):
     """
     Build messages from the conversations and images.
@@ -83,28 +84,38 @@ def preprocess(
 
     # get texts from conversations
     converstion = get_conv_template(sft_format)
-    # only use the last 3 round of messages
-    latest_messages = messages[-3:]
+    # only use the last 10 round of messages
+    latest_messages = messages[-10:]
+    results.append(
+        {
+                    "role": "system",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": override_system_prompt if override_system_prompt else converstion.system_message,
+                        }
+                    ],
+        }
+    )
+    print("The actual system prompt for generation:", override_system_prompt if override_system_prompt else converstion.system_message)
     for mid, message in enumerate(latest_messages):
         if message["role"] == converstion.roles[0] or message["role"] == "user":
             record = {
                 "role": message["role"],
                 "content": [],
             }
-            if "images" in message:
+            if "timestamps" in message and "images" in message and message["timestamps"] is not None:
+                per_round_images, per_round_timestamps = message["images"], message["timestamps"]
+                for image, timestamp in zip(per_round_images, per_round_timestamps):
+                    images.append(image)
+                    record["content"].append({"type": "text", "text": f"{int(timestamp)//3600:02d}:{(int(timestamp)//60-60*(int(timestamp)//3600)):02d}:{int(timestamp)%60:02d}"})
+                    record["content"].append({"type": "image", "image": image})                    
+            elif "images" in message:
                 per_round_images = message["images"]
-                if len(per_round_images) > 2:
-                    per_round_images = per_round_images[-2:]
-                    print(f"Only use the last 2 images in the {mid}-th round")
-
-                images.extend(per_round_images)
                 for image in per_round_images:
-                    record["content"].append(
-                        {
-                            "type": "image",
-                            "image": image,
-                        }
-                    )
+                    images.append(image)
+                    record["content"].append({"type": "image", "image": image})
+                    
             if 'content' in message:
                 record["content"].append(
                     {
@@ -113,6 +124,7 @@ def preprocess(
                     }
                 )
             results.append(record)
+            
         elif message["role"] == converstion.roles[1] or message["role"] == "assistant":
             formatted_answer = message["content"].strip()
             # ◁think▷用户说了“你好”，这是一个非常简单的问候，通常用于开启对话。我需要判断用户的意图。可能性一：用户只是礼貌性地打招呼，想要开启一段对话；可能性二：用户可能有更具体的需求，比如询问我的功能、功能或者需要帮助。由于用户没有提供更多信息，我需要保持开放，同时引导用户进一步说明他们的需求。
@@ -137,7 +149,7 @@ def preprocess(
                 formatted_answer.count(processor.image_token) == 0
             ), f"there should be no {processor.image_token} in the assistant's reply, but got {messages}"
             converstion.append_message(converstion.roles[1], formatted_answer)
-
+    
     text = processor.apply_chat_template(results, add_generation_prompt=True)
     print(f"raw text = {text}")
     if len(images) == 0:
@@ -153,11 +165,13 @@ def preprocess(
     return inputs
 
 
-
+@torch.no_grad()
+@torch.inference_mode()
 def kimi_vl_generate(
     model: torch.nn.Module,
     processor: AutoProcessor,
     conversations: list[Conversation],
+    override_system_prompt,
     stop_words: list,
     max_length: int = 256,
     temperature: float = 1.0,
@@ -166,7 +180,7 @@ def kimi_vl_generate(
 ):
     # convert conversation to inputs
     print(f"conversations = {conversations}")
-    inputs = preprocess(conversations, processor=processor)
+    inputs = preprocess(conversations, processor=processor, override_system_prompt=override_system_prompt)
     inputs = inputs.to(model.device)
 
     return generate(
@@ -180,7 +194,6 @@ def kimi_vl_generate(
         chunk_size=chunk_size,
     )
 
-
 def generate(
     model,
     processor,