import os
import requests
from uuid import uuid4
from PIL import Image
import gradio as gr

def gen_image(prompt: str):
    """generate the image from the chinese stable diffusion model of paddlenlp server

    Args:
        prompt (str): the source of the prompt
    """
    if not prompt:
        return
    access_token = os.environ['token']
    url = f"https://aip.baidubce.com/rpc/2.0/nlp-itec/poc/stable_diffusion?access_token={access_token}&text={prompt}"
    content = requests.get(url).content
    cache_dir = 'images'
    os.makedirs(cache_dir, exist_ok=True)

    tempfile = os.path.join(cache_dir, f'{str(uuid4())}.png')
    with open(tempfile, 'wb') as f:
        f.write(content)
    image = Image.open(tempfile)
    os.remove(tempfile)
    return [image]

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

block = gr.Blocks(css=read_content('assets/style.css'))

examples = [
    '贝尼·赖特森、丹·蒙福德、亚伦·霍尔基的黑白血腥维多利亚小镇夜景特写街景，恐怖，月亮升起，交叉阴影，高对比度，超精细，极简主义构图，4k',
    '马头雕塑插图印刷，超精细，丹·蒙福德，亚伦·霍基，高对比度，低聚合风格',
    '一只熊从篝火旁的冰箱里偷食物，黑白雕刻版画，交叉影线',
    '厚涂层油画《悲伤男孩的特写肖像》，本·科迪，希卡里·希莫达'
]

with block:
    gr.HTML(read_content("assets/header.html"))
    gr.Markdown("[![Stargazers repo roster for @PaddlePaddle/PaddleNLP](https://reporoster.com/stars/PaddlePaddle/PaddleNLP)](https://github.com/PaddlePaddle/PaddleNLP)")
    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="输入中文，生成图片",
                ).style(
                    border=(True, False, True, True),
                    rounded=(True, False, False, True),
                    container=False,
                )

                btn = gr.Button("Generate image").style(
                    margin=False,
                    rounded=(False, True, True, False),
                )

        gallery = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery"
        ).style(grid=[1, 1], height="auto")

        gr.Examples(examples=examples, fn=gen_image, inputs=text, outputs=gallery)
        text.submit(gen_image, inputs=text, outputs=gallery)
        btn.click(gen_image, inputs=text, outputs=gallery)

    gr.Image('./assets/paddlenlp-preview.jpeg')
    gr.HTML(read_content("assets/footer.html"))

block.queue(concurrency_count=5).launch(debug=True)