update APIs

app.py

@@ -11,7 +11,6 @@ from css_and_js import js, call_JS
 from PIL import Image, PngImagePlugin, ImageChops
 
 url_host = "http://flagstudio.baai.ac.cn"
-#token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoiZjAxOGMxMzJiYTUyNDBjMzk5NTMzYTI5YjBmMzZiODMiLCJhcHBfbmFtZSI6IndlYiIsImlkZW50aXR5X3R5cGUiOiIyIiwidXNlcl9yb2xlIjoiMiIsImp0aSI6IjVjMmQzMjdiLWI5Y2MtNDhiZS1hZWQ4LTllMjQ4MDk4NzMxYyIsIm5iZiI6MTY2OTAwNjE5NywiZXhwIjoxOTg0MzY2MTk3LCJpYXQiOjE2NjkwMDYxOTd9.9B3MDk8wA6iWH5puXjcD19tJJ4Ox7mdpRyWZs5Kwt70"
 token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoiMGY4M2QxMDg3N2MzMTFlZGFiYzYwZmU5ZGFjMTI1ZDMiLCJhcHBfbmFtZSI6IndlYiIsImlkZW50aXR5X3R5cGUiOiIyIiwidXNlcl9yb2xlIjoiMiIsImp0aSI6ImE3YTE1N2I3LTllNTItNDllMS04YzA0LWEzZmI5YjZiZjNlYSIsIm5iZiI6MTY3MDU5MTcwMSwiZXhwIjoxOTg1OTUxNzAxLCJpYXQiOjE2NzA1OTE3MDF9.OcfGayna-wr_5mo4LT6OJHSCokna8vqKSmmCftFUsx8"
 
 def read_content(file_path: str) -> str:
@@ -58,7 +57,6 @@ def upload_image(img):
 def post_reqest(seed, prompt, width, height, image_num, img=None, mask=None):
     data = {
             "type": "gen-image",
-            "gen_image_num": image_num,
             "parameters": {
                 "width": width, # output height width
                 "height": height, # output image height
@@ -87,35 +85,47 @@ def post_reqest(seed, prompt, width, height, image_num, img=None, mask=None):
                         "height": mask.height,
                     }
     headers = {"token": token}
+
     # Send create task request
-    # url = "http://flagstudio.baai.ac.cn/api/v1/task/create"
+    all_task_data = []
     url = url_host+"/api/v1/task/create"
-    r = requests.post(url, json=data, headers=headers)
-    if r.status_code != 200:
-        raise gr.Error(r.reason)
-    create_res = r.json()
-    task_id = create_res["data"]["task_id"]
+    for _ in range(image_num):
+        r = requests.post(url, json=data, headers=headers)
+        if r.status_code != 200:
+            raise gr.Error(r.reason)
+        create_res = r.json()
+        if create_res['code'] == 3002:
+            raise gr.Error("Inappropriate prompt detected.")
+        elif create_res['code'] != 0:
+            raise gr.Error("Unknown error")
+        all_task_data.append(create_res["data"])
 
     # Get result
     url = url_host+"/api/v1/task/status"
+    images = []
     while True:
-        r = requests.post(url, json=create_res["data"], headers=headers)
-        if r.status_code != 200:
-            raise gr.Error(r.reason)
-        res = r.json()
-        if res["code"] == 6002:
-            # Running
-            time.sleep(1)
-            continue
-        elif res["code"] == 0:
-            # Finished
-            images = []
-            for img_info in res["data"]["images"]:
-                img_res = requests.get(img_info["url"])
-                images.append(Image.open(io.BytesIO(img_res.content)).convert("RGB"))
+        if len(all_task_data) <= 0:
             return images
-        else:
-            raise gr.Error(f"Error code: {res['code']}")
+        for i in range(len(all_task_data)-1, -1, -1):
+            data = all_task_data[i]
+            r = requests.post(url, json=data, headers=headers)
+            if r.status_code != 200:
+                raise gr.Error(r.reason)
+            res = r.json()
+            if res["code"] == 6002:
+                # Running
+                continue
+            if res["code"] == 6005:
+                raise gr.Error("NSFW image detected.")
+            elif res["code"] == 0:
+                # Finished
+                for img_info in res["data"]["images"]:
+                    img_res = requests.get(img_info["url"])
+                    images.append(Image.open(io.BytesIO(img_res.content)).convert("RGB"))
+                del all_task_data[i]
+            else:
+                raise gr.Error(f"Error code: {res['code']}")
+        time.sleep(1)
 
 def request_images(raw_text, class_draw, style_draw, batch_size, w, h, seed):
     if filter_content(class_draw) != "国画":
@@ -134,7 +144,13 @@ def request_images(raw_text, class_draw, style_draw, batch_size, w, h, seed):
 
 
 def img2img(prompt, image_and_mask):
-    return post_reqest(0, prompt, 512, 512, 1, image_and_mask["image"], image_and_mask["mask"])
+    if image_and_mask["image"].width <= image_and_mask["image"].height:
+        width = 512
+        height = int((width/image_and_mask["image"].width)*image_and_mask["image"].height)
+    else:
+        height = 512
+        width = int((height/image_and_mask["image"].height)*image_and_mask["image"].width)
+    return post_reqest(0, prompt, width, height, 1, image_and_mask["image"], image_and_mask["mask"])
 
 
 examples = [
@@ -311,4 +327,4 @@ if __name__ == "__main__":
         gr.HTML(read_content("footer.html"))
         # gr.Image('./contributors.png')
 
-    block.queue(max_size=50, concurrency_count=20).launch()
+    block.queue(max_size=100, concurrency_count=50).launch()