import gradio as gr

from tiny_image_gen.catalog import STYLE_PRESETS, default_prompt, style_choices
from tiny_image_gen.service import TinyImageService


service = TinyImageService()


def update_style(style_name: str):
    return STYLE_PRESETS[style_name].hint


def run_generation(prompt: str, style_name: str, negative_prompt: str, steps: int, guidance: float, seed: int):
    return service.generate(
        prompt=prompt,
        style_name=style_name,
        negative_prompt=negative_prompt,
        steps=steps,
        guidance=guidance,
        seed=seed,
    )


with gr.Blocks(title="Tiny Text To Image CPU") as demo:
    gr.Markdown(
        """
        # Tiny Text To Image CPU
        Small text-to-image generation running on a free CPU Space.

        - Model: `segmind/tiny-sd`
        - Separate Space
        - CPU-friendly defaults
        - Single-image generation
        """
    )

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                value=default_prompt(),
                lines=6,
            )
            style = gr.Dropdown(
                label="Style",
                choices=style_choices(),
                value="Cinematic",
            )
            style_hint = gr.Textbox(
                label="Style Hint",
                value=STYLE_PRESETS["Cinematic"].hint,
                interactive=False,
                lines=3,
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                value="blurry, low quality, distorted, deformed, extra fingers, watermark, text",
                lines=3,
            )
            steps = gr.Slider(
                label="Steps",
                minimum=4,
                maximum=20,
                value=10,
                step=1,
            )
            guidance = gr.Slider(
                label="Guidance Scale",
                minimum=1.0,
                maximum=10.0,
                value=6.0,
                step=0.5,
            )
            seed = gr.Number(
                label="Seed",
                value=42,
                precision=0,
            )
            generate = gr.Button("Generate Image", variant="primary")

        with gr.Column():
            image = gr.Image(label="Image", type="pil")
            status = gr.Textbox(label="Status", value=service.describe())
            info = gr.Textbox(
                label="Info",
                value="This Space uses a compact diffusion model, so quality is lower than large GPU models but it fits free CPU hardware better.",
                lines=6,
            )

    style.change(
        fn=update_style,
        inputs=style,
        outputs=style_hint,
    )

    generate.click(
        fn=run_generation,
        inputs=[prompt, style, negative_prompt, steps, guidance, seed],
        outputs=[image, status, info],
    )


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