import torch
from diffusers import ShapEPipeline
import gradio as gr
from diffusers.utils import export_to_gif
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def display_gif(gif_path):
    return gif_path

pipe = ShapEPipeline.from_pretrained("openai/shap-e", torch_dtype=torch.float16, variant="fp16")
pipe = pipe.to(device)

guidance_scale = 15.0
prompt = ["A firecracker", "A birthday cupcake"]

images = pipe(
    prompt,
    guidance_scale=guidance_scale,
    num_inference_steps=64,
    frame_size=256,
).images


gif_file = export_to_gif(images[0], "firecracker_3d.gif")
export_to_gif(images[1], "cake_3d.gif")




interface = gr.Interface(fn=display_gif, inputs=[], outputs="image")

interface.launch(share=True, examples=[gif_file])