| import os |
| from typing import Dict, Any |
| from PIL import Image |
| from io import BytesIO |
|
|
| from inference import 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()} |
|
|
