app.py

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 = "https://flagstudio.baai.ac.cn"
token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoiMGY4M2QxMDg3N2MzMTFlZGFiYzYwZmU5ZGFjMTI1ZDMiLCJhcHBfbmFtZSI6IndlYiIsImlkZW50aXR5X3R5cGUiOiIyIiwidXNlcl9yb2xlIjoiMiIsImp0aSI6ImE3YTE1N2I3LTllNTItNDllMS04YzA0LWEzZmI5YjZiZjNlYSIsIm5iZiI6MTY3MDU5MTcwMSwiZXhwIjoxOTg1OTUxNzAxLCJpYXQiOjE2NzA1OTE3MDF9.OcfGayna-wr_5mo4LT6OJHSCokna8vqKSmmCftFUsx8"

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",
            "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
    all_task_data = []
    url = url_host+"/api/v1/task/create"
    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:
        if len(all_task_data) <= 0:
            return images
        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) != "国画":
        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):
    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 = [
    '水墨蝴蝶和牡丹花，国画',
    '苍劲有力的墨竹，国画',
    '暴风雨中的灯塔',
    '机械小松鼠，科学幻想',
    '中国水墨山水画，国画',
    "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=512, concurrency_count=256).launch()