Update app.py

app.py

@@ -1,50 +1,36 @@
 import os
 import requests
 import json
-import random
 import gradio as gr
 from PIL import Image
 from io import BytesIO
 
-from utils import get_token
 
-
-def deliver_request(style, desc):
-    number = random.randint(1, 90) % 3
-    print("node: ", number)
-    return generate_figure(style, desc, number)
-
-
-def generate_figure(style, desc, number):
-    token = get_token()
-    url_node1 = os.environ["url_node1"]
-    url_node2 = os.environ["url_node2"]
-    url_node3 = os.environ["url_node3"]
-    url_list = [url_node1, url_node2, url_node3]
+def generate_figure(style, desc):
+    url = os.environ["req_url"]
 
     requests_json = {
-        "user_name": "huggingface",
         "style": style,
         "desc": desc
     }
 
     headers = {
         "Content-Type": "application/json",
-        "X-Auth-Token": token
     }
 
-    response = requests.post(url_list[number], json=requests_json, headers=headers, verify=False)
+    response = requests.post(url, json=requests_json, headers=headers, verify=False)
+
+    status = response.status_code
+    assert status == 201
+
     response = json.loads(response.text)
-    status = response["status"]
-    assert status == 200
-    url_list = response["output_image_url"]
 
-    image1 = Image.open(BytesIO(requests.get(url_list[0]).content))
-    image2 = Image.open(BytesIO(requests.get(url_list[1]).content))
-    image3 = Image.open(BytesIO(requests.get(url_list[2]).content))
-    image4 = Image.open(BytesIO(requests.get(url_list[3]).content))
+    url_dict = response["data"]["pictures"]
+    image_list = []
+    for k in url_dict:
+        image_list.append(Image.open(BytesIO(requests.get(url_dict[k]).content)))
 
-    return image1, image2, image3, image4
+    return image_list
 
 
 def read_content(file_path: str) -> str:
@@ -69,7 +55,7 @@ css = """
 """
 
 # warm up
-generate_figure("梵高", "一只猫", 0)
+generate_figure("梵高", "一只猫")
 
 with gr.Blocks(css=css) as demo:
     gr.HTML(read_content("./header.html"))
@@ -116,7 +102,7 @@ with gr.Blocks(css=css) as demo:
         gr.Markdown("- Try [Wukong-Huahua model on the Foundation Models Platform for Mindspore]"
                     "(https://xihe.mindspore.cn/modelzoo/wukong)")
 
-    generate_button.click(deliver_request,
+    generate_button.click(generate_figure,
                           inputs=[style_input, desc_input],
                           outputs=[img_output1, img_output2, img_output3, img_output4])