import gradio as gr
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("prompthero/openjourney") # about 750 s
# pipe = DiffusionPipeline.from_pretrained("admruul/anything-v3.0") # about 2500 s
# pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2") # about 1700s, poor quality?
# pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") # about 850s

# prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
# image = pipe(prompt).images[0]

def generate(input_prompt):
    generations = pipe(input_prompt)
    generated_img = generations.images[0]
    return generated_img

gradio_app = gr.Interface(
    generate,
    inputs=gr.Text(label="What kind of image would you like to generate?"),
    outputs=[gr.Image(label="Generated Image")],
    title="IMAGE GENERATION!!",
)

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