Update app.py

app.py

@@ -8,47 +8,45 @@ def generate_3d_model(prompt, output_path="assistant_3d.obj"):
     Generate a 3D model using ShapE and export it in a Blender-compatible format
     """
     try:
-        # Load pipeline with memory constraints suitable for CPU-only usage
+        # Load pipeline optimized for CPU usage
         pipe = ShapEPipeline.from_pretrained(
             "openai/shap-e",
             torch_dtype=torch.float32,
-            low_cpu_mem_usage=True  # Minimize CPU memory usage
+            low_cpu_mem_usage=True
         ).to("cpu")
         
-        # Generate with reduced inference steps and guidance scale for CPU
+        # Generate model
         outputs = pipe(
             prompt=prompt,
-            num_inference_steps=8,  # Reduced steps for CPU speed
-            guidance_scale=5.0,     # Lower guidance scale to save memory
+            num_inference_steps=8,  # Reduced for CPU efficiency
+            guidance_scale=5.0      # Lower guidance to conserve resources
         )
         
-        # Extract vertices and faces for the 3D mesh generation
-        vertices = outputs["vertices"][0].detach().cpu().numpy()
-        faces = outputs["faces"][0].detach().cpu().numpy()
+        # Check output structure for debugging
+        print("Output structure:", outputs.keys())
         
-        # Create the trimesh object with error handling
-        mesh_obj = trimesh.Trimesh(vertices=vertices, faces=faces, process=True)
-        
-        # Export model in desired format with robust handling for compatibility
+        # Attempt to extract vertices and faces
         try:
-            if output_path.endswith('.obj'):
-                mesh_obj.export(output_path, include_normals=True)
-            else:
-                mesh_obj.export(output_path)
-            print(f"Successfully exported 3D model to: {output_path}")
-        except Exception as export_error:
-            print(f"Error during export: {export_error}")
-            # Use a fallback format in case of export failure
-            backup_path = output_path.rsplit('.', 1)[0] + '.ply'
-            mesh_obj.export(backup_path)
-            print(f"Exported backup model to: {backup_path}")
+            vertices = outputs["vertices"][0].detach().cpu().numpy()
+            faces = outputs["faces"][0].detach().cpu().numpy()
+        except KeyError as ke:
+            print(f"Key error: {ke}")
+            print(f"Available keys in output: {list(outputs.keys())}")
+            return None
+        
+        # Construct the 3D mesh object
+        mesh_obj = trimesh.Trimesh(vertices=vertices, faces=faces, process=True)
         
+        # Export mesh object
+        if output_path.endswith('.obj'):
+            mesh_obj.export(output_path, include_normals=True)
+        else:
+            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}")
-        print(f"Error type: {type(e)}")
-        print(f"Full error details: {str(e)}")
         raise
 
 if __name__ == "__main__":