Update model.py

model.py

@@ -1,24 +1,21 @@
+# model.py
 import torch
 from diffusers import ShapEPipeline
 from diffusers.utils import export_to_gif
 
-# Define checkpoint ID and load pipeline on CPU
-ckpt_id = "openai/shap-e"
-pipe = ShapEPipeline.from_pretrained(ckpt_id).to("cpu")
+# 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
 
-# Define generation parameters
-guidance_scale = 10.0  # Lowered for efficiency on CPU
-num_inference_steps = 32  # Reduced steps for CPU performance
-prompt = "a shark"
-
-# Generate images from the prompt with reduced settings
-images = pipe(
-    prompt=prompt,
-    guidance_scale=guidance_scale,
-    num_inference_steps=num_inference_steps,
-    size=256,  # Image size for the model
-).images
-
-# Export images to GIF format
-gif_path = export_to_gif(images, "shark_3d.gif")
-print(f"GIF saved at {gif_path}")
+# 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,
+        size=size
+    ).images
+    gif_path = export_to_gif(images, "generated_3d.gif")
+    return gif_path