new

handler.py

@@ -1,52 +1,7 @@
-import os
-from typing import Dict, Any
-from PIL import Image
-from io import BytesIO
+# handler.py
 
-from inference import Chat  # 直接import你放的inference.py里Chat类
-from robohusky.conversation import get_conv_template
+from inference import Chat
+chat = Chat()
 
-class EndpointHandler:
-    def __init__(self, path: str = "."):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.chat = Chat(
-            model_path=path,
-            device=self.device,
-            num_gpus=1,
-            max_new_tokens=1024,
-            load_8bit=False
-        )
-        self.vision_feature = None
-        self.modal_type = "text"
-        self.conv = get_conv_template("husky").copy()
-
-    def preprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
-        query = inputs.get("inputs", "")
-        self.conv = get_conv_template("husky").copy()
-        self.vision_feature = None
-        self.modal_type = "text"
-
-        if "image" in inputs:
-            image_bytes = inputs["image"]
-            image = Image.open(BytesIO(image_bytes)).convert("RGB")
-            image.save("temp.jpg")
-            self.vision_feature = self.chat.get_image_embedding("temp.jpg")
-            self.modal_type = "image"
-
-        elif "video" in inputs:
-            video_bytes = inputs["video"]
-            with open("temp.mp4", "wb") as f:
-                f.write(video_bytes)
-            self.vision_feature = self.chat.get_video_embedding("temp.mp4")
-            self.modal_type = "video"
-
-        return {"query": query}
-
-    def __call__(self, inputs: Dict[str, Any]) -> Dict[str, str]:
-        processed = self.preprocess(inputs)
-        query = processed["query"]
-
-        conversations = self.chat.ask(text=query, conv=self.conv, modal_type=self.modal_type)
-        outputs = self.chat.answer(conversations, self.vision_feature, modal_type=self.modal_type)
-        self.conv.messages[-1][1] = outputs.strip()
-        return {"output": outputs.strip()}
+def inference_fn(inputs):
+    return chat.answer(inputs)