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")
pipe.enable_xformers_memory_efficient_attention()


# Function to generate an image from text using diffusion
def generate_image(prompt):
    
    images = pipe(prompt=prompt)
    return images[0]

# Gradio interface
inputs = gr.inputs.Textbox(lines=5, label="Enter text to generate image:")
outputs = gr.outputs.Image(label="Generated Image")

title = "ShoeGen: Generate an Image of a Shoe"
description = "Enter a text description of a shoe to generate an image of the shoe."
examples = [["A red shoe with white laces and black sole."], ["A blue sneaker with a white stripe."], ["A brown boot with a buckle."]]

gr.Interface(
    fn=generate_image,
    inputs=inputs,
    outputs=outputs,
    title=title,
    description=description,
    examples=examples
).launch()