File size: 686 Bytes
fe4aa3a 487d3dd fe4aa3a 487d3dd fe4aa3a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
# model.py
import torch
from diffusers import ShapEPipeline
from diffusers.utils import export_to_gif
# Load pipeline once to avoid reloading with each request
def load_pipeline():
ckpt_id = "openai/shap-e"
pipe = ShapEPipeline.from_pretrained(ckpt_id, torch_dtype=torch.float32, trust_remote_code=True).to("cpu")
return pipe
# Generate images and export to GIF
def generate_3d_gif(pipe, prompt, guidance_scale=10.0, num_inference_steps=32, size=256):
images = pipe(
prompt=prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
).images
gif_path = export_to_gif(images, "generated_3d.gif")
return gif_path
|