import os
import random

import gradio as gr
from gradio_client import Client, file
import numpy as np
from PIL import Image
from typing import Tuple
from translatepy import Translator

MODEL = os.environ.get("MODEL")
API_URL = "https://api-inference.huggingface.co/models/tianweiy/DMD2"
DESCRIPTION = """
# DMD2 文生图
"""
translator = Translator()



def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

MAX_SEED = np.iinfo(np.int32).max

client = Client(MODEL)

style_list = [
    {
        "name": "(无风格)",
        "prompt": "{prompt}",
    },
    {
        "name": "电影",
        "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
    },
    {
        "name": "摄影",
        "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed",
    },
    {
        "name": "动画",
        "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed",
    },
    {
        "name": "漫画",
        "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style",
    },
    {
        "name": "数绘",
        "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
    },
    {
        "name": "像素",
        "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
    },
    {
        "name": "幻想",
        "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
    },
    {
        "name": "朋克",
        "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
    },
    {
        "name": "三维",
        "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
    },
]   
styles = {k["name"]: (k["prompt"]) for k in style_list}
print(styles)
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "(无风格)"

def apply_style(style_name: str, positive: str) -> Tuple[str, str]:
    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    return p.replace("{prompt}", positive), n

def generate(
    prompt: str,
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    style: str = DEFAULT_STYLE_NAME,
    num_images: int = 2,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    prompt = str(translator.translate(prompt, 'English'))
    print(prompt)
    seed = int(randomize_seed_fn(seed, randomize_seed))
    # print(client.view_api())
    result = client.predict(
		prompt=prompt,
		seed=seed,
		height=height,
		width=width,
		num_images=num_images,
		fast_vae_decode=True,
		api_name="/inference"
    )
    images = result[0]
    print(images)
    image_paths = []
  #  List[Dict(image: filepath, caption: str | None)]
    for img in images:
        image_paths.append(img["image"])
    print(image_paths)
    return image_paths, seed

examples = [
    "镭射眼的秋田犬",
    "一只吃起司的猫",
    "太空中骑马的宇航员",
    "放学回家的学生们，动画风格",
    "一个可爱的机器人艺术家在画架上绘画，概念艺术",
    "一位女士的特写，她戴着透明、棱柱形、精致的复仇女神头饰，摆出应有的姿势，棕色肤色"
]

CSS = '''
.gradio-container{max-width: 560px !important}
h1{text-align:center}
footer {
    visibility: hidden
}
'''
with gr.Blocks(css=CSS, theme="soft") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Group():
        with gr.Row():
            prompt = gr.Text(
                label="描述",
                show_label=False,
                max_lines=1,
                placeholder="画什么好呢",
                container=False,
                scale=2,
            )
            run_button = gr.Button("生成", scale=1)
        result = gr.Gallery(label="作品", columns=1, preview=True)
    with gr.Accordion("高级选项", open=False):
        with gr.Row():
            num_images = gr.Slider(
                label="数量",
                minimum=1,
                maximum=5,
                step=1,
                value=2,
            )
            seed = gr.Slider(
                label="种子",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
                visible=True
            )
            randomize_seed = gr.Checkbox(label="随机种子", value=True)
        with gr.Row(visible=True):
            width = gr.Slider(
                label="宽",
                minimum=512,
                maximum=2048,
                step=8,
                value=1024,
            )
            height = gr.Slider(
                label="高",
                minimum=512,
                maximum=2048,
                step=8,
                value=1024,
            )
    with gr.Row(visible=True):
            style_selection = gr.Radio(
                show_label=True,
                container=True,
                interactive=True,
                choices=STYLE_NAMES,
                value=DEFAULT_STYLE_NAME,
                label="风格化",
            )
        

    gr.Examples(
        examples=examples,
        inputs=prompt,
        outputs=[result, seed],
        fn=generate,
        cache_examples="lazy",
    )
    
    gr.on(
        triggers=[
            prompt.submit,
            run_button.click,
        ],
        fn=generate,
        inputs=[
            prompt,
            seed,
            width,
            height,
            style_selection,
            num_images,
            randomize_seed,
        ],
        outputs=[result, seed],
        api_name="run",
    )


        
if __name__ == "__main__":
    demo.queue(max_size=20).launch(show_api=False, debug=False)