import gradio as gr
from diffusers import DiffusionPipeline
import torch

# Load the pipeline
pipeline = DiffusionPipeline.from_pretrained("openai/shap-e", torch_dtype=torch.float16)
pipeline = pipeline.to("cuda" if torch.cuda.is_available() else "cpu")

def generate_image(prompt):
    # Generate image using the prompt
    output = pipeline(prompt)
    return output.images[0]

# Set up Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="Shap-E Image Generation",
    description="Generate images from text prompts using the Shap-E diffusion model."
)

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