import gradio as gr

from config import (
    APP_DESCRIPTION,
    APP_TITLE,
    DEFAULT_GUIDANCE,
    DEFAULT_HEIGHT,
    DEFAULT_NEGATIVE_PROMPT,
    DEFAULT_NUM_IMAGES,
    DEFAULT_NUM_STEPS,
    DEFAULT_PROMPT,
    DEFAULT_SEED,
    DEFAULT_WIDTH,
    EXAMPLE_PROMPTS,
    MAX_GUIDANCE,
    MAX_HEIGHT,
    MAX_NUM_IMAGES,
    MAX_NUM_STEPS,
    MAX_WIDTH,
    MIN_GUIDANCE,
    MIN_NUM_STEPS,
)
from models import run_generation
from utils import prepare_generator, sanitize_dimensions


def generate_images(
    prompt: str,
    negative_prompt: str,
    guidance_scale: float,
    num_inference_steps: int,
    width: int,
    height: int,
    num_images: int,
    seed: int,
):
    if not prompt.strip():
        raise gr.Error("Please provide a prompt to describe your image.")
    width, height = sanitize_dimensions(width, height, MAX_WIDTH, MAX_HEIGHT)
    generator, final_seed = prepare_generator(seed)
    images = run_generation(
        prompt=prompt.strip(),
        negative_prompt=negative_prompt.strip() if negative_prompt else None,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        num_images=num_images,
        generator=generator,
    )
    return images, final_seed


with gr.Blocks(fill_width=True) as demo:
    gr.Markdown(
        f"""
# {APP_TITLE}

{APP_DESCRIPTION}

[Built with anycoder](https://huggingface.co/spaces/akhaliq/anycoder)
"""
    )

    with gr.Row():
        prompt_input = gr.Textbox(
            label="Prompt",
            value=DEFAULT_PROMPT,
            placeholder="Describe what you want to see...",
            lines=3,
        )
        negative_input = gr.Textbox(
            label="Negative Prompt",
            value=DEFAULT_NEGATIVE_PROMPT,
            placeholder="Specify what to avoid (optional)",
            lines=3,
        )

    with gr.Accordion("Generation Settings", open=False):
        with gr.Row():
            guidance_slider = gr.Slider(
                minimum=MIN_GUIDANCE,
                maximum=MAX_GUIDANCE,
                step=0.1,
                value=DEFAULT_GUIDANCE,
                label="Guidance Scale",
            )
            steps_slider = gr.Slider(
                minimum=MIN_NUM_STEPS,
                maximum=MAX_NUM_STEPS,
                step=1,
                value=DEFAULT_NUM_STEPS,
                label="Inference Steps",
            )
        with gr.Row():
            width_slider = gr.Slider(
                minimum=256,
                maximum=MAX_WIDTH,
                step=8,
                value=DEFAULT_WIDTH,
                label="Width (px)",
            )
            height_slider = gr.Slider(
                minimum=256,
                maximum=MAX_HEIGHT,
                step=8,
                value=DEFAULT_HEIGHT,
                label="Height (px)",
            )
        with gr.Row():
            num_images_slider = gr.Slider(
                minimum=1,
                maximum=MAX_NUM_IMAGES,
                step=1,
                value=DEFAULT_NUM_IMAGES,
                label="Images per prompt",
            )
            seed_number = gr.Number(
                value=DEFAULT_SEED,
                label="Seed (-1 for random)",
                precision=0,
            )

    generate_button = gr.Button("Generate", variant="primary")
    gallery = gr.Gallery(
        label="Generated Images",
        columns=2,
        height="auto",
        object_fit="contain",
        show_share_button=True,
    )
    seed_display = gr.Number(
        label="Used Seed",
        value=DEFAULT_SEED,
        interactive=False,
        precision=0,
    )

    generate_button.click(
        fn=generate_images,
        inputs=[
            prompt_input,
            negative_input,
            guidance_slider,
            steps_slider,
            width_slider,
            height_slider,
            num_images_slider,
            seed_number,
        ],
        outputs=[gallery, seed_display],
        api_name="generate",
    )

    gr.Examples(
        examples=[[example] for example in EXAMPLE_PROMPTS],
        inputs=[prompt_input],
        label="Prompt Ideas",
    )

    gr.ClearButton(
        components=[prompt_input, negative_input, gallery],
        value="Reset",
    )

demo.queue(max_size=40).launch()