#!/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)
demo.queue(max_size=10).launch()