moooji commited on
Commit
2cf0517
·
1 Parent(s): f36f1a2

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +22 -0
handler.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoImageProcessor, Swinv2Model
2
+ import torch
3
+
4
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
5
+
6
+ class EndpointHandler():
7
+ def __init__(self, path=""):
8
+ self.model = Swinv2Model.from_pretrained("microsoft/moooji/swinv2-large-patch4-window12to24-192to384-22kto1k-ft").to(device)
9
+ self.processor = AutoImageProcessor.from_pretrained("microsoft/moooji/swinv2-large-patch4-window12to24-192to384-22kto1k-ft")
10
+
11
+ def __call__(self, data: Any) -> List[float]:
12
+ inputs = data.pop("inputs", data)
13
+
14
+ image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
15
+ inputs = self.processor(image, return_tensors="pt").to(device)
16
+
17
+ with torch.no_grad():
18
+ outputs = model(**inputs)
19
+
20
+ last_hidden_states = outputs.last_hidden_state
21
+ return last_hidden_states[2].tolist()
22
+