SlowPacer commited on
Commit
52c421d
1 Parent(s): f52fb63

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +4 -2
handler.py CHANGED
@@ -35,8 +35,10 @@ class EndpointHandler():
35
  inputs = data.pop("inputs", data)
36
  parameters = data.pop("parameters", {})
37
 
38
- inputs = isinstance(inputs, list) and inputs or [inputs]
39
- raw_images = [Image.open(BytesIO(base64.b64decode(_img))) for _img in inputs]
 
 
40
 
41
  processed_images = self.processor(images=raw_images, return_tensors="pt")
42
  processed_images["pixel_values"] = processed_images["pixel_values"].to(device)
 
35
  inputs = data.pop("inputs", data)
36
  parameters = data.pop("parameters", {})
37
 
38
+ if isinstance(inputs, Image):
39
+ raw_images = [inputs]
40
+ else:
41
+ raw_images = [Image.open(BytesIO(base64.b64decode(_img))) for _img in inputs]
42
 
43
  processed_images = self.processor(images=raw_images, return_tensors="pt")
44
  processed_images["pixel_values"] = processed_images["pixel_values"].to(device)