# -*- coding: utf-8 -*-
# app.py
# Simple Gradio Space with Stable Diffusion (diffusers)
# Designed for fast generation and medium quality

import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

MODEL_ID = "stabilityai/sd-turbo"
DEFAULT_WIDTH = 512
DEFAULT_HEIGHT = 512
DEFAULT_STEPS = 20
DEFAULT_SCALE = 7.5


def load_pipeline(model_id=MODEL_ID):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        revision="fp16" if device == "cuda" else None,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    )

    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to(device)

    try:
        if device == "cuda":
            pipe.enable_attention_slicing()
            pipe.enable_xformers_memory_efficient_attention()
    except Exception:
        pass

    return pipe


pipeline = load_pipeline()


def generate_image(prompt, negative_prompt, width, height, steps, guidance_scale, seed):
    if not prompt:
        return None

    generator = None
    if seed is not None and seed >= 0:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        generator = torch.Generator(device).manual_seed(seed)

    max_pixels = 768 * 768
    if width * height > max_pixels:
        scale = (max_pixels / (width * height)) ** 0.5
        width = int(width * scale)
        height = int(height * scale)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        with torch.autocast("cuda"):
            image = pipeline(
                prompt=prompt,
                negative_prompt=negative_prompt if negative_prompt else None,
                height=height,
                width=width,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                generator=generator,
            ).images[0]
    else:
        image = pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt else None,
            height=height,
            width=width,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            generator=generator,
        ).images[0]

    return image


with gr.Blocks(title="Fast — Medium Quality Stable Diffusion") as demo:
    gr.Markdown("## Simple Stable Diffusion Image Generator")

    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(label="Prompt", placeholder="a portrait of an astronaut riding a horse")
            negative = gr.Textbox(label="Negative prompt", placeholder="low quality, blurry")
            with gr.Row():
                width_in = gr.Number(value=DEFAULT_WIDTH, label="Width")
                height_in = gr.Number(value=DEFAULT_HEIGHT, label="Height")
            with gr.Row():
                steps_in = gr.Slider(minimum=5, maximum=50, value=DEFAULT_STEPS, step=1, label="Steps")
                scale_in = gr.Slider(minimum=1.0, maximum=12.0, value=DEFAULT_SCALE, step=0.1, label="Guidance Scale")
            seed_in = gr.Number(value=-1, label="Seed (-1 for random)")
            run = gr.Button("Generate")

        with gr.Column(scale=1):
            output = gr.Image(label="Output", type="pil")

    run.click(
        fn=generate_image,
        inputs=[prompt, negative, width_in, height_in, steps_in, scale_in, seed_in],
        outputs=[output],
    )

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