import gradio as gr
from stable_diffusion_tf.stable_diffusion import StableDiffusion
from PIL import Image

generator = StableDiffusion(
    img_height=512,
    img_width=512,
    jit_compile=True,
)


def gen(prompt):
    image = generator.generate(
        prompt,
        num_steps=50,
        unconditional_guidance_scale=7.5,
        temperature=1,
        batch_size=1,
    )
    return image[0]

demo = gr.Interface(fn=gen, inputs="text", outputs=gr.Image(type="pil"))
    
demo.launch()

# prompt = "a photo of an astronaut riding a horse on mars"
# image = pipe(prompt).images[0]  
# image.save("astronaut_rides_horse.png")