Add application file

app.py

@@ -0,0 +1,318 @@
+import io 
+import re
+import imp
+import time
+import json 
+import base64
+import requests
+import gradio as gr
+import ui_functions as uifn
+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"
+
+def read_content(file_path: str) -> str:
+    """read the content of target file
+    """
+    with open(file_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+
+    return content
+
+def filter_content(raw_style: str):
+    if "(" in raw_style:
+        i = raw_style.index("(")
+    else :
+        i = -1
+
+    if i == -1:
+        return raw_style
+    else :
+        return raw_style[:i]
+
+def upload_image(img):
+    url = url_host + "/api/v1/image/get-upload-link"
+    headers = {"token": token}
+    r = requests.post(url, json={}, headers=headers)
+    if r.status_code != 200:
+        raise gr.Error(r.reason)
+    head_res = r.json()
+    if head_res["code"] != 0:
+        raise gr.Error("Unknown error")
+    image_id = head_res["data"]["image_id"]
+    image_url = head_res["data"]["url"]
+    image_headers = head_res["data"]["headers"]
+
+    imgBytes = io.BytesIO()
+    img.save(imgBytes, "PNG")
+    imgBytes = imgBytes.getvalue()
+
+    r = requests.put(image_url, data=imgBytes, headers=image_headers)
+    if r.status_code != 200:
+        raise gr.Error(r.reason)
+    return image_id, image_url
+
+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
+                "prompts": [prompt], 
+            }
+        }
+    data["parameters"]["seed"] = int(seed)
+    if img is not None:
+        # Upload image
+        image_id, image_url = upload_image(img)
+        data["parameters"]["init_image"] = {
+                    "image_id": image_id,
+                    "url": image_url,
+                    "width": img.width,
+                    "height": img.height,
+                }
+    if mask is not None:
+        # Upload mask
+        extrama = mask.convert("L").getextrema()
+        if extrama[1] > 0:
+            mask_id, mask_url = upload_image(mask)
+            data["parameters"]["mask_image"] = {
+                        "image_id": mask_id,
+                        "url": mask_url,
+                        "width": mask.width,
+                        "height": mask.height,
+                    }
+    headers = {"token": token}
+    # Send create task request
+    # url = "http://flagstudio.baai.ac.cn/api/v1/task/create"
+    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"]
+
+    # Get result
+    url = url_host+"/api/v1/task/status"
+    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"))
+            return images
+        else:
+            raise gr.Error(f"Error code: {res['code']}")
+
+def request_images(raw_text, class_draw, style_draw, batch_size, w, h, seed):
+    if filter_content(class_draw) != "国画":
+        if filter_content(class_draw) != "通用":
+            raw_text = raw_text + f"，{filter_content(class_draw)}"
+
+        for sty in style_draw:
+            raw_text = raw_text + f"，{filter_content(sty)}"
+    elif filter_content(class_draw) == "国画":
+        raw_text = raw_text + "，国画，水墨画，大作，黑白，高清，传统"
+    print(f"raw text is {raw_text}")
+
+    images = post_reqest(seed, raw_text, w, h, int(batch_size))
+
+    return images
+
+
+def img2img(prompt, image_and_mask):
+    return post_reqest(0, prompt, 512, 512, 1, image_and_mask["image"], image_and_mask["mask"])
+
+
+examples = [
+    '水墨蝴蝶和牡丹花，国画',
+    '苍劲有力的墨竹，国画',
+    '暴风雨中的灯塔',
+    '机械小松鼠，科学幻想',
+    '中国水墨山水画，国画',
+    "Lighthouse in the storm",
+    "A dog",
+    "Landscape by 张大千",
+    "A tiger 长��兔子耳朵",
+    "A baby bird 铅笔素描",
+]
+
+if __name__ == "__main__":
+    block = gr.Blocks(css=read_content('style.css'))
+
+    with block:
+        gr.HTML(read_content("header.html"))
+        with gr.Tabs(elem_id='tabss') as tabs:
+
+            with gr.TabItem("文生图(Text-to-img)", id='txt2img_tab'):
+
+                with gr.Group():
+                    with gr.Box():
+                        with gr.Row().style(mobile_collapse=False, equal_height=True):
+                            text = gr.Textbox(
+                                label="Prompt",
+                                show_label=False,
+                                max_lines=1,
+                                placeholder="Input text(输入文字)",
+                                interactive=True,
+                            ).style(
+                                border=(True, False, True, True),
+                                rounded=(True, False, False, True),
+                                container=False,
+                            )
+                            
+                            btn = gr.Button("Generate image").style(
+                                margin=False,
+                                rounded=(True, True, True, True),
+                            )
+                        with gr.Row().style(mobile_collapse=False, equal_height=True):
+                            class_draw = gr.Radio(choices=["通用(general)","国画(traditional Chinese painting)",], value="通用(general)", show_label=True, label='生成类型(type)')
+                            # class_draw = gr.Dropdown(["通用(general)", "国画(traditional Chinese painting)", 
+                            #                             "照片，摄影(picture photography)", "油画(oil painting)", 
+                            #                             "铅笔素描(pencil sketch)", "CG", 
+                            #                             "水彩画(watercolor painting)", "水墨画(ink and wash)",
+                            #                             "插画(illustrations)", "3D", "图生图(img2img)"],
+                            #                         label="生成类型(type)",
+                            #                         show_label=True,
+                            #                         value="通用(general)")
+                        with gr.Row().style(mobile_collapse=False, equal_height=True):
+                            style_draw = gr.CheckboxGroup(["蒸汽朋克(steampunk)", "电影摄影风格(film photography)", 
+                                                        "概念艺术(concept art)", "Warming lighting", 
+                                                        "Dramatic lighting", "Natural lighting", 
+                                                        "虚幻引擎(unreal engine)", "4k", "8k",
+                                                        "充满细节(full details)"],
+                                                    label="画面风格(style)",
+                                                    show_label=True,
+                                                    )
+                        with gr.Row().style(mobile_collapse=False, equal_height=True):
+                            # sample_size = gr.Slider(minimum=1,
+                            #                         maximum=4,
+                            #                         step=1,
+                            #                         label="生成数量(number)",
+                            #                         show_label=True,
+                            #                         interactive=True,
+                            #                         )
+                            sample_size = gr.Radio(choices=["1","2","3","4"], value="1", show_label=True, label='生成数量(number)')
+                            seed = gr.Number(0, label='seed', interactive=True)
+                        with gr.Row().style(mobile_collapse=False, equal_height=True):
+                            w = gr.Slider(512,1024,value=512, step=64, label="width")
+                            h = gr.Slider(512,1024,value=512, step=64, label="height")
+
+                    gallery = gr.Gallery(
+                        label="Generated images", show_label=False, elem_id="gallery"
+                    ).style(grid=[2,2])
+                    gr.Examples(examples=examples, fn=request_images, inputs=text, outputs=gallery, examples_per_page=100)
+                    with gr.Row().style(mobile_collapse=False, equal_height=True):
+                        img_choices = gr.Dropdown(["图片1(img1)"],label='请选择一张图片发送到图生图',show_label=True,value="图片1(img1)")
+                    with gr.Row().style(mobile_collapse=False, equal_height=True):
+                        output_txt2img_copy_to_input_btn = gr.Button("发送图片到图生图(Sent the image to img2img)").style(
+                                margin=False,
+                                rounded=(True, True, True, True),
+                            )
+
+                    with gr.Row():
+                        prompt = gr.Markdown("提示(Prompt)：", visible=False)
+                    with gr.Row():
+                        move_prompt_zh = gr.Markdown("请移至图生图部分进行编辑(拉到顶部)", visible=False)
+                    with gr.Row():
+                        move_prompt_en = gr.Markdown("Please move to the img2img section for editing(Pull to the top)", visible=False)
+
+                        
+
+                    text.submit(request_images, inputs=[text, class_draw, style_draw, sample_size, w, h, seed], outputs=gallery)
+                    btn.click(request_images, inputs=[text, class_draw, style_draw, sample_size, w, h, seed], outputs=gallery)
+
+                    sample_size.change(
+                        fn=uifn.change_img_choices,
+                        inputs=[sample_size],
+                        outputs=[img_choices]
+                    )
+
+            with gr.TabItem("图生图(Img-to-Img)", id="img2img_tab"):
+                with gr.Row(elem_id="prompt_row"):
+                    img2img_prompt = gr.Textbox(label="Prompt",
+                                                elem_id='img2img_prompt_input',
+                                                placeholder="神奇的森林，流淌的河流.",
+                                                lines=1,
+                                                max_lines=1,
+                                                value="",
+                                                show_label=False).style()
+
+                    img2img_btn_mask = gr.Button("Generate", variant="primary", visible=False,
+                                                 elem_id="img2img_mask_btn")
+                    img2img_btn_editor = gr.Button("Generate", variant="primary", elem_id="img2img_edit_btn")
+                gr.Markdown('#### 输入图像')
+                with gr.Row().style(equal_height=False):
+                    #with gr.Column():
+                    img2img_image_mask = gr.Image(
+                        value=None,
+                        source="upload",
+                        interactive=True,
+                        tool="sketch",
+                        type='pil',
+                        elem_id="img2img_mask",
+                        image_mode="RGBA"
+                    )
+                gr.Markdown('#### 编辑后的图片')
+                with gr.Row():
+                    output_img2img_gallery = gr.Gallery(label="Images", elem_id="img2img_gallery_output").style(
+                        grid=[4,4,4] )
+                with gr.Row():
+                    gr.Markdown('提示(prompt):')
+                with gr.Row():
+                    gr.Markdown('请选择一张图像掩盖掉一部分区域,并输入文本描述')
+                with gr.Row():
+                    gr.Markdown('Please select an image to cover up a part of the area and enter a text description.')
+                gr.Markdown('# 编辑设置',visible=False)
+
+
+                output_txt2img_copy_to_input_btn.click(
+                    uifn.copy_img_to_input,
+                    [gallery, img_choices],
+                    [tabs, img2img_image_mask, move_prompt_zh, move_prompt_en, prompt]
+                )
+
+
+                img2img_func = img2img
+                img2img_inputs = [img2img_prompt, img2img_image_mask]
+                img2img_outputs = [output_img2img_gallery]
+
+                img2img_btn_mask.click(
+                    img2img_func,
+                    img2img_inputs,
+                    img2img_outputs
+                )
+
+                def img2img_submit_params():
+                    return (img2img_func,
+                            img2img_inputs,
+                            img2img_outputs)
+
+                img2img_btn_editor.click(*img2img_submit_params())
+
+                # GENERATE ON ENTER
+                img2img_prompt.submit(None, None, None,
+                                      _js=call_JS("clickFirstVisibleButton",
+                                                  rowId="prompt_row"))
+
+        gr.HTML(read_content("footer.html"))
+        # gr.Image('./contributors.png')
+
+    block.queue(max_size=50, concurrency_count=20).launch(
+            # share=True,
+            show_error=True,
+            server_name="0.0.0.0",
+            server_port=43523,
+            )