import gradio as gr
import torch
from diffusers import ShapEPipeline
from diffusers.utils import export_to_gif

# Load the ShapE model
ckpt_id = "openai/shap-e"
pipe = ShapEPipeline.from_pretrained(ckpt_id)

def generate_shap_e_gif(prompt):
    guidance_scale = 15.0
    images = pipe(
        prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=64,
    ).images

    gif_path = export_to_gif(images, f"{prompt}_3d.gif")
    return gif_path

# Create the Gradio interface
demo = gr.Interface(
    fn=generate_shap_e_gif,
    inputs=gr.Textbox(lines=2, placeholder="Enter a prompt"),
    outputs=gr.File(),
    title="ShapE 3D GIF Generator",
    description="Enter a prompt to generate a 3D GIF using the ShapE model."
)

# Run the app
if __name__ == "__main__":
    demo.launch()