import gradio as gr
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, progress=gr.Progress()):
    guidance_scale = 15.0
    num_inference_steps = 64
    progress(0, desc="Starting...")
    images = []
    
    for i in range(num_inference_steps):
        image = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=1).images[0]
        images.append(image)
        # Update the progress tracker
        progress((i+1)/num_inference_steps)
    
    gif_path = export_to_gif(images, f"{prompt}_3d.gif")
    # Ensure the progress is set to complete
    progress(1, desc="Completed")
    return gif_path

# Create the Gradio interface with queue enabled
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."
).queue()

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