add sr choice

app.py

@@ -7,6 +7,7 @@ import requests
 import json 
 import ui_functions as uifn
 from css_and_js import js, call_JS
+import re
 
 txt2img_defaults = {
     'prompt': '',
@@ -120,7 +121,26 @@ def request_images(raw_text, class_draw, style_draw, batch_size, sr_option):
         # print(content[i])
         images.append(base2picture(content[i]))        
 
-    return images 
+    return images
+
+def sr_request_images(img_str, idx):
+    idx_map = {
+            "图片1(img1)":0,
+            "图片2(img2)":1,
+            "图片3(img3)":2,
+            "图片4(img4)":3,
+        }
+    idx = idx_map[idx]
+    image_data = re.sub('^data:image/.+;base64,', '', img_str[idx])
+    d = {"data":[image_data, 0]} # batch_size设置为0，
+    url = "http://flagart.baai.ac.cn/api/general/"
+    r = requests.post(url, json=d, headers={"Content-Type": "application/json", "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br", "Connection": "keep-alive"})
+    result_text = r.text
+    content = json.loads(result_text)["data"][0]
+    images = [base2picture(content[0])] 
+
+    return images
+
 
 def encode_pil_to_base64(pil_image):
     with io.BytesIO() as output_bytes:
@@ -247,11 +267,20 @@ if __name__ == "__main__":
                     ).style(grid=[2], height="auto")
 
                     with gr.Row().style(mobile_collapse=False, equal_height=True):
-                        img_choices = gr.Dropdown(["图片1(img1)"],label='请选择一张图片发送到图生图',show_label=True,value="图片1(img1)")
+                        img_choices = gr.Dropdown(["图片1(img1)"],label='请选择一张图片发送到图生图或者超分',show_label=True,value="图片1(img1)")
                         output_txt2img_copy_to_input_btn = gr.Button("发送图片到图生图(Sent the image to img2img)").style(
                                 margin=False,
                                 rounded=(True, True, True, True),
                             )
+                        output_txt2img_sr_btn = gr.Button("将选择的图片进行超分").style(
+                                margin=False,
+                                rounded=(True, True, True, True),
+                            )
+                    sr_gallery = gr.Gallery(
+                        label="SR images", show_label=True, elem_id="sr_gallery"
+                    ).style(grid=[2], height="auto")
+
+                        
 
                     # 并没有什么软用，转跳还是很卡，而且只能转跳一次
                     # go_to_img2img_btn = gr.Button("转至图生图(Go to img2img)").style(
@@ -263,6 +292,8 @@ if __name__ == "__main__":
                     gr.Examples(examples=examples, fn=request_images, inputs=text, outputs=gallery, examples_per_page=100)
                     text.submit(request_images, inputs=[text, class_draw, style_draw, sample_size, sr_option], outputs=gallery)
                     btn.click(request_images, inputs=[text, class_draw, style_draw, sample_size, sr_option], outputs=gallery)
+                    # to do here
+                    output_txt2img_sr_btn.click(sr_request_images, inputs=[gallery, img_choices], outputs=[sr_gallery])
 
                     sample_size.change(
                         fn=uifn.change_img_choices,