Update app.py

app.py

@@ -1,38 +1,48 @@
 import torch
 from diffusers import ShapEPipeline
-from diffusers.utils import export_to_gif
-import PIL.Image
+import trimesh
+import numpy as np
 
-def generate_3d_model(prompt, output_path="assistant_3d.gif"):
+def generate_3d_model(prompt, output_path="assistant_3d.obj"):
     """
-    Generate a 3D model using ShapE optimized for CPU usage
+    Generate a 3D model using ShapE and export it in a Blender-compatible format
     """
     try:
-        # Force CPU and reduced precision
+        # Load pipeline with minimal settings
         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
+        # Generate with minimal parameters
         outputs = pipe(
             prompt,
-            num_inference_steps=32,  # Reduced from default
-            frame_size=32,  # Smaller frame size
-            guidance_scale=10.0,  # Reduced guidance scale
+            num_inference_steps=16,
+            frame_size=24,
+            guidance_scale=7.5
         )
         
-        # 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
+        # Get the mesh data
+        mesh = outputs.meshes[0]  # Get the first mesh
+        
+        # Convert to vertices and faces format
+        verts = mesh.verts.detach().cpu().numpy()
+        faces = mesh.faces.detach().cpu().numpy()
+        
+        # Create trimesh object
+        mesh_obj = trimesh.Trimesh(vertices=verts, faces=faces)
+        
+        # Export in desired format
+        if output_path.endswith('.obj'):
+            mesh_obj.export(output_path)
+        elif output_path.endswith('.glb'):
+            mesh_obj.export(output_path)
+        elif output_path.endswith('.stl'):
+            mesh_obj.export(output_path)
+        
+        print(f"Successfully exported 3D model to: {output_path}")
+        return output_path
         
     except Exception as e:
         print(f"Error during generation: {e}")
@@ -41,8 +51,9 @@ def generate_3d_model(prompt, output_path="assistant_3d.gif"):
         raise
 
 if __name__ == "__main__":
-    prompt = "A gentle AI voice assistant constructed from a circle ring and 3 lines that fly alongside the circle"  # Simplified prompt
+    prompt = "a simple ring"
     try:
-        generate_3d_model(prompt)
+        # You can change the extension to .glb or .stl as needed
+        generate_3d_model(prompt, "assistant_3d.obj")
     except Exception as e:
         print(f"Generation failed: {e}")