import gradio as gr
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline
import random

model = "hakurei/waifu-diffusion"
device = "cpu"

pipe = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch.float32)
pipe = pipe.to(device)

block = gr.Blocks(css=".container { max-width: 800px; margin: auto; }")

def infer(prompt, width, height, nums, steps, guidance_scale, seed):
    print(prompt)
    print(width, height, nums, steps, guidance_scale, seed)

    if prompt is not None and prompt != "":
        if seed is None or seed == '' or seed == -1:
            seed = int(random.randrange(4294967294))
        generator = torch.Generator(device).manual_seed(seed)
        images = pipe([prompt] * nums, height=height, width=width, num_inference_steps=steps, generator=generator, guidance_scale=guidance_scale )["sample"]
    return images

# with block as demo:
def run():
    _app = gr.Interface(
        fn=infer,
        title="Waifu Generator",
        inputs=[
            gr.Textbox(label="prompt"),
            gr.Slider(512, 1024, 512, step=64, label="width"),
            gr.Slider(512, 1024, 512, step=64, label="height"),
            gr.Slider(1, 4, 1, step=1, label="Number of Images"),
            gr.Slider(10, 150, step=1, value=50,
                      label="num_inference_steps:\n"
                            "The number of denoising steps. More de-scaling steps usually result in a higher quality image, but will slow down inference."),
            gr.Slider(0, 20, 7.5, step=0.5,
                      label="guidance_scale:\n" +
                            "A higher boot ratio encourages the generation of images that are closely related to text \"hints\", often at the expense of reduced image quality"),
            gr.Textbox(label="Random seed",
                       placeholder="Random Seed",
                       lines=1),
        ],
        outputs=[
            gr.Gallery(label="Generated images")
        ])

    return _app

app = run()
app.launch(debug=True)