File size: 1,912 Bytes
cf932d8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import os
from typing import Dict, Any
from PIL import Image
from io import BytesIO
from inference import Chat # 直接import你放的inference.py里Chat类
from robohusky.conversation import get_conv_template
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()}
|