rbanfield commited on
Commit
3c0b02a
·
1 Parent(s): e36f852

Code is working

Browse files
Files changed (1) hide show
  1. handler.py +0 -8
handler.py CHANGED
@@ -1,4 +1,3 @@
1
- import sys
2
  from io import BytesIO
3
  import base64
4
 
@@ -6,8 +5,6 @@ from PIL import Image
6
  import torch
7
  from transformers import CLIPProcessor, CLIPTextModel, CLIPVisionModelWithProjection
8
 
9
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
10
-
11
  class EndpointHandler():
12
  def __init__(self, path=""):
13
  self.text_model = CLIPTextModel.from_pretrained("rbanfield/clip-vit-large-patch14")
@@ -16,19 +13,14 @@ class EndpointHandler():
16
 
17
  def __call__(self, data):
18
  inputs = data.pop("inputs", None)
19
- print(inputs, file=sys.stderr)
20
  text_input = inputs["text"] if "text" in inputs else None
21
  image_input = inputs["image"] if "image" in inputs else None
22
 
23
  if text_input:
24
- print("in text mode", file=sys.stderr)
25
- print(text_input, file=sys.stderr)
26
  processor = self.processor(text=text_input, return_tensors="pt", padding=True)
27
  with torch.no_grad():
28
  return self.text_model(**processor).pooler_output.tolist()
29
  elif image_input:
30
- print("in image mode", file=sys.stderr)
31
- print(image_input, file=sys.stderr)
32
  image = Image.open(BytesIO(base64.b64decode(image_input)))
33
  processor = self.processor(images=image, return_tensors="pt")
34
  with torch.no_grad():
 
 
1
  from io import BytesIO
2
  import base64
3
 
 
5
  import torch
6
  from transformers import CLIPProcessor, CLIPTextModel, CLIPVisionModelWithProjection
7
 
 
 
8
  class EndpointHandler():
9
  def __init__(self, path=""):
10
  self.text_model = CLIPTextModel.from_pretrained("rbanfield/clip-vit-large-patch14")
 
13
 
14
  def __call__(self, data):
15
  inputs = data.pop("inputs", None)
 
16
  text_input = inputs["text"] if "text" in inputs else None
17
  image_input = inputs["image"] if "image" in inputs else None
18
 
19
  if text_input:
 
 
20
  processor = self.processor(text=text_input, return_tensors="pt", padding=True)
21
  with torch.no_grad():
22
  return self.text_model(**processor).pooler_output.tolist()
23
  elif image_input:
 
 
24
  image = Image.open(BytesIO(base64.b64decode(image_input)))
25
  processor = self.processor(images=image, return_tensors="pt")
26
  with torch.no_grad():