+import spaces
import gradio as gr
import torch
from diffusers import DiffusionPipeline

model_id = "Apocalypse-19/shoe-generator"
pipe = DiffusionPipeline.from_pretrained(model_id)

pipe.to("cuda")

# Function to generate an image from text using diffusion
+@spaces.GPU
def generate_image(prompt):
    
    images = pipe(prompt=prompt, negative_prompt=neg_prompt,)
    return images[0]

_TITLE = "Shoe Generator"
with gr.Blocks(_TITLE) as ShoeGen:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(lines=3, label="Enter a prompt")
            neg_prompt = gr.Textbox(lines=3, label="Enter a negative prompt")
        with gr.Column():
            button_gen = gr.Button("Generate Image")
            output = gr.Image(label="Generated Image")

    button_gen.click(generate_image, inputs=[prompt, neg_prompt], outputs=output)

ShoeGen.launch()