#!/usr/bin/env python

from __future__ import annotations

import gradio as gr

from model import Model

DESCRIPTION = "# [StyleGAN-Human](https://github.com/stylegan-human/StyleGAN-Human)"

model = Model()

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            with gr.Row():
                seed1 = gr.Number(label="Seed 1", value=6876)
                psi1 = gr.Slider(label="Truncation psi 1", minimum=0, maximum=2, step=0.05, value=0.7)
            with gr.Row():
                generate_button1 = gr.Button("Generate")
            with gr.Row():
                generated_image1 = gr.Image(label="Generated Image 1", type="numpy", height=600)

        with gr.Column():
            with gr.Row():
                seed2 = gr.Number(label="Seed 2", value=6886)
                psi2 = gr.Slider(label="Truncation psi 2", minimum=0, maximum=2, step=0.05, value=0.7)
            with gr.Row():
                generate_button2 = gr.Button("Generate")
            with gr.Row():
                generated_image2 = gr.Image(label="Generated Image 2", type="numpy", height=600)

    with gr.Row():
        with gr.Column():
            with gr.Row():
                num_frames = gr.Slider(label="Number of Intermediate Frames", minimum=0, maximum=41, step=1, value=7)
            with gr.Row():
                interpolate_button = gr.Button("Interpolate")
            with gr.Row():
                interpolated_images = gr.Gallery(label="Output Images", object_fit="scale-down")

    generate_button1.click(
        model.generate_single_image,
        inputs=[
            seed1,
            psi1,
        ],
        outputs=generated_image1,
    )
    generate_button2.click(
        model.generate_single_image,
        inputs=[
            seed2,
            psi2,
        ],
        outputs=generated_image2,
    )
    interpolate_button.click(
        model.generate_interpolated_images,
        inputs=[
            seed1,
            psi1,
            seed2,
            psi2,
            num_frames,
        ],
        outputs=interpolated_images,
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch()