from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("prompthero/openjourney")

# 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()