File size: 1,685 Bytes
f164923 |
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 |
# handler.py
import torch
from transformers import pipeline
# check for GPU
device = 0 if torch.cuda.is_available() else -1
# multi-model list
multi_model_list = [
{"id": "a", "model_id": "bytedance-research/UI-TARS-72B-DPO", "task": " image-text-to-text"},
{"id": "b", "model_id": "bytedance-research/UI-TARS-72B-DPO", "task": " image-text-to-text"},
{"id": "c", "model_id": "bytedance-research/UI-TARS-72B-DPO", "task": " image-text-to-text"},
{"id": "d", "model_id": "bytedance-research/UI-TARS-72B-DPO", "task": " image-text-to-text"},
{"id": "e", "model_id": "bytedance-research/UI-TARS-72B-DPO", "task": " image-text-to-text"},
]
class EndpointHandler():
def __init__(self, path=""):
self.multi_model={}
# load all the models onto device
for model in multi_model_list:
self.multi_model[model["model_id"]] = pipeline(model["task"], model=model["model_id"], device=device)
def __call__(self, data):
# deserialize incomin request
inputs = data.pop("inputs", data)
parameters = data.pop("parameters", None)
model_id = data.pop("model_id", None)
# check if model_id is in the list of models
if model_id is None or model_id not in self.multi_model:
raise ValueError(f"model_id: {model_id} is not valid. Available models are: {list(self.multi_model.keys())}")
# pass inputs with all kwargs in data
if parameters is not None:
prediction = self.multi_model[model_id](inputs, **parameters)
else:
prediction = self.multi_model[model_id](inputs)
# postprocess the prediction
return prediction |