Quoc Bao Bui commited on
Commit
57ccf6b
1 Parent(s): 7ef0124

Add local request script

Browse files
Files changed (1) hide show
  1. local_request.py +33 -0
local_request.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import base64
3
+ from io import BytesIO
4
+
5
+ from PIL import Image
6
+
7
+ from handler import EndpointHandler, decode_base64_image
8
+
9
+
10
+ def local_predict(prompts, encode_image):
11
+ # Init handler
12
+ my_handler = EndpointHandler()
13
+ if encode_image:
14
+ response = my_handler({"inputs": prompts, "image": encode_image})
15
+ else:
16
+ response = my_handler({"inputs": prompts})
17
+
18
+ image = decode_base64_image(response["image"])
19
+ image.save("local_output.png")
20
+
21
+
22
+ opt = argparse.ArgumentParser("Diffuser local test")
23
+ opt.add_argument("-prompts", "--prompts", default="", type=str, help="Diffuser prompts")
24
+ opt.add_argument("-image", "--image", default="", type=str, help="Init image")
25
+ if __name__ == '__main__':
26
+ args = opt.parse_args()
27
+
28
+ encoded_string = ""
29
+ if args.image:
30
+ with open(args.image, "rb") as image_file:
31
+ encoded_string = base64.b64encode(image_file.read()).decode()
32
+
33
+ local_predict(args.prompts, encoded_string)