jeff-RQ commited on
Commit
6b0c845
·
1 Parent(s): 32f759e

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +4 -1
handler.py CHANGED
@@ -1,6 +1,7 @@
1
  from typing import Any, Dict
2
  from transformers import Blip2Processor, Blip2ForConditionalGeneration
3
-
 
4
 
5
  class EndpointHandler:
6
  def __init__(self, path=""):
@@ -16,6 +17,8 @@ class EndpointHandler:
16
  image = data.pop("image", data)
17
  text = data.pop("text", data)
18
 
 
 
19
  inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device)
20
  generated_ids = self.model.generate(**inputs)
21
  generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
 
1
  from typing import Any, Dict
2
  from transformers import Blip2Processor, Blip2ForConditionalGeneration
3
+ import io
4
+ from PIL import Image
5
 
6
  class EndpointHandler:
7
  def __init__(self, path=""):
 
17
  image = data.pop("image", data)
18
  text = data.pop("text", data)
19
 
20
+ image = Image.open(io.BytesIO(image))
21
+
22
  inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device)
23
  generated_ids = self.model.generate(**inputs)
24
  generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()