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

# get hf user access token as an environment variable
TOKEN_KEY = os.getenv('AUTH_TOKEN')

# setup pipeline
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=TOKEN_KEY)
pipe = pipe.to('cuda')

# define gradio function
def generate(prompt:str, seed:int, guidance:float):
    generator = torch.Generator("cuda").manual_seed(int(seed))
    with autocast("cuda"):
        image = pipe(prompt=prompt, generator=generator, guidance_scale=guidance, steps=50).images[0]
    return image

# create the gradio UI
demo = gr.Interface(
    fn=generate,
    inputs=[gr.Textbox(placeholder="castle on a mountain"), gr.Number(value=123456), gr.Slider(0,10)],
    outputs="image",
    allow_flagging="never",
)

# allow queueing or incoming requests, max=3 
demo.queue(concurrency_count=3)

# launch demo
demo.launch()