metafresh89 commited on
Commit
f65a75f
1 Parent(s): 272f751

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +26 -0
handler.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from PIL import Image
3
+ from io import BytesIO
4
+ from transformers import pipeline
5
+ import base64
6
+
7
+ class EndpointHandler():
8
+ def __init__(self, path=""):
9
+ self.pipeline=pipeline("controlnet_qrcode-control_v11p_sd21",model=path)
10
+
11
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
12
+ """
13
+ data args:
14
+ images (:obj:`string`)
15
+ candiates (:obj:`list`)
16
+ Return:
17
+ A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX", "score": 0.82}
18
+ """
19
+ inputs = data.pop("inputs", data)
20
+ prompts = data.pop("prompt", data)
21
+
22
+ image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
23
+
24
+ # run prediction one image wit provided candiates
25
+ prediction = self.pipeline(images=[image], prompt=[prompts])
26
+ return prediction[0]