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