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

import re
import os
import requests

def greet(name):
    return "Hello " + name + "!!"

def create_waifu(prompt):
    pipe = StableDiffusionPipeline.from_pretrained(
    'hakurei/waifu-diffusion',
    torch_dtype=torch.float16
    )

    image = pipe(prompt, guidance_scale=6)[0][0]

    return image
"""
Interface class is initialized with three required parameters
- fn: the function to wrap a UI around
- inputs: components to use for the input
- outputs: which components to use for the output
"""
iface = gr.Interface(fn=create_waifu, 
                    inputs=gr.Textbox(lines=1, placeholder="Please enter your prompt..."), 
                    outputs="image")
iface.launch()