Update handler.py
Browse files- handler.py +5 -4
handler.py
CHANGED
@@ -7,21 +7,22 @@ from transformers import AutoProcessor, LlavaForConditionalGeneration
|
|
7 |
class EndpointHandler():
|
8 |
def __init__(self, path=""):
|
9 |
model_id = ""
|
10 |
-
model = LlavaForConditionalGeneration.from_pretrained(
|
11 |
model_id,
|
12 |
torch_dtype=torch.float16,
|
13 |
low_cpu_mem_usage=True,
|
14 |
).to(0)
|
15 |
-
processor = AutoProcessor.from_pretrained(model_id)
|
16 |
|
17 |
def __call__(self, data: Dict[str, Any]):
|
18 |
parameters = data.pop("inputs",data)
|
19 |
inputs = data.pop("inputs", data)
|
20 |
if parameters is not None:
|
21 |
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
|
22 |
-
|
|
|
23 |
prompt = "USER: <image>\nWhat are these?\nASSISTANT:"
|
24 |
-
output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
|
25 |
return output
|
26 |
|
27 |
|
|
|
7 |
class EndpointHandler():
|
8 |
def __init__(self, path=""):
|
9 |
model_id = ""
|
10 |
+
self.model = LlavaForConditionalGeneration.from_pretrained(
|
11 |
model_id,
|
12 |
torch_dtype=torch.float16,
|
13 |
low_cpu_mem_usage=True,
|
14 |
).to(0)
|
15 |
+
self.processor = AutoProcessor.from_pretrained(model_id)
|
16 |
|
17 |
def __call__(self, data: Dict[str, Any]):
|
18 |
parameters = data.pop("inputs",data)
|
19 |
inputs = data.pop("inputs", data)
|
20 |
if parameters is not None:
|
21 |
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
|
22 |
+
raw_image = Image.open(requests.get(url, stream=True).raw)
|
23 |
+
inputs = self.processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16)
|
24 |
prompt = "USER: <image>\nWhat are these?\nASSISTANT:"
|
25 |
+
output = self.model.generate(**inputs, max_new_tokens=200, do_sample=False)
|
26 |
return output
|
27 |
|
28 |
|