Include pointcloud2mesh components to additionally view 3d model

app.py

@@ -9,6 +9,9 @@ from point_e.diffusion.sampler import PointCloudSampler
 from point_e.models.download import load_checkpoint
 from point_e.models.configs import MODEL_CONFIGS, model_from_config
 from point_e.util.plotting import plot_point_cloud
+from point_e.util.pc_to_mesh import marching_cubes_mesh
+
+import trimesh
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -29,6 +32,14 @@ base_model.load_state_dict(load_checkpoint(base_name, device))
 print('downloading upsampler checkpoint...')
 upsampler_model.load_state_dict(load_checkpoint('upsample', device))
 
+print('creating SDF model...')
+name = 'sdf'
+sdf_model = model_from_config(MODEL_CONFIGS[name], device)
+sdf_model.eval()
+
+print('loading SDF model...')
+sdf_model.load_state_dict(load_checkpoint(name, device))
+
 sampler = PointCloudSampler(
     device=device,
     models=[base_model, upsampler_model],
@@ -65,12 +76,30 @@ def inference(prompt):
             )
         ),
     )
-    return fig
+    
+    # Produce a mesh (with vertex colors)
+    mesh = marching_cubes_mesh(
+        pc=pc,
+        model=sdf_model,
+        batch_size=4096,
+        grid_size=32, # increase to 128 for resolution used in evals
+        progress=True,
+    )
+
+    # Write the mesh to a PLY file to import into some other program.
+    with open("mesh.ply", 'wb') as f:
+        mesh.write_ply(f)
+    
+    obj_file = '3d_model.obj'
+    mesh = trimesh.load('mesh.ply')
+    mesh.export(obj_file)
+
+    return fig, obj_file
 
 demo = gr.Interface(
     fn=inference,
     inputs="text",
-    outputs=gr.Plot(),
+    outputs=[gr.Plot(),gr.Model3D(value=None)],
     examples=[
         ["a red motorcycle"],
         ["a RED pumpkin"],