Update app.py

app.py

@@ -1,24 +1,49 @@
 import torch
 from diffusers import ShapEPipeline
 from diffusers.utils import export_to_gif
+import PIL.Image
 
-ckpt_id = "openai/shap-e"
-device = torch.device("cpu")
+def generate_3d_model(prompt, output_path="assistant_3d.gif"):
+    """
+    Generate a 3D model using ShapE optimized for CPU usage
+    """
+    try:
+        # Force CPU and reduced precision
+        pipe = ShapEPipeline.from_pretrained(
+            "openai/shap-e",
+            torch_dtype=torch.float32,
+            low_cpu_mem_usage=True
+        ).to("cpu")
+        
+        # Minimal generation settings to reduce memory usage
+        outputs = pipe(
+            prompt,
+            num_inference_steps=32,  # Reduced from default
+            frame_size=32,  # Smaller frame size
+            guidance_scale=10.0,  # Reduced guidance scale
+            num_frames=30  # Reduced number of frames
+        )
+        
+        # Ensure we have PIL images
+        if not isinstance(outputs.images[0], PIL.Image.Image):
+            images = [PIL.Image.fromarray(img) for img in outputs.images]
+        else:
+            images = outputs.images
+            
+        # Save as GIF
+        gif_path = export_to_gif(images, output_path)
+        print(f"Successfully created GIF at: {gif_path}")
+        return gif_path
+        
+    except Exception as e:
+        print(f"Error during generation: {e}")
+        print(f"Error type: {type(e)}")
+        print(f"Full error details: {str(e)}")
+        raise
 
-# Attempt to load the pipeline without auth for testing
-try:
-    pipe = ShapEPipeline.from_pretrained(ckpt_id).to(device)
-except Exception as e:
-    print(f"Error loading model: {e}")
-    exit(1)
-
-guidance_scale = 15.0
-prompt = "A gentle AI voice assistant constructed from a circle ring and 3 lines that fly alongside the circle"
-
-# Run the model with minimal parameters to troubleshoot
-try:
-    images = pipe(prompt, guidance_scale=guidance_scale).images
-    gif_path = export_to_gif(images, "assistant_3d.gif")
-    print(f"GIF created at: {gif_path}")
-except Exception as e:
-    print(f"Error generating images: {e}")
+if __name__ == "__main__":
+    prompt = "A gentle AI voice assistant constructed from a circle ring and 3 lines that fly alongside the circle"  # Simplified prompt
+    try:
+        generate_3d_model(prompt)
+    except Exception as e:
+        print(f"Generation failed: {e}")