import gradio as gr
from inference import generate

default_prompt = """Cinematic portrait of a Japanese girl as Shinobu Kocho, realistic facial features, gradient purple compound eyes, intricate butterfly hair ornament, wearing silk haori with butterfly wing patterns, soft moonlight, hyper-realistic, depth of field, purple butterfly particles, ethereal lighting"""


def run(prompt, height, width, steps, seed):
    return generate(
        prompt,
        height,
        width,
        steps,
        seed,
    )


with gr.Blocks() as demo:
    gr.Markdown("## Z-Image Turbo")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(value=default_prompt, lines=4)

            height = gr.Slider(256, 1024, value=512, step=64)
            width = gr.Slider(256, 1024, value=512, step=64)

            steps = gr.Slider(1, 20, value=9, step=1)
            seed = gr.Number(value=-1)

            btn = gr.Button("Generate")

        with gr.Column():
            output = gr.Image()

    btn.click(
        fn=run,
        inputs=[prompt, height, width, steps, seed],
        outputs=output,
    )


if __name__ == "__main__":
    demo.launch()