Code is working
Browse files- 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():
|