Update handler.py

handler.py

@@ -8,10 +8,13 @@ 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
+from point_e.util.point_cloud import PointCloud
 import json
 import base64
 import numpy as np
 from io import BytesIO
+import os
 
 
 # set device
@@ -24,6 +27,7 @@ class EndpointHandler():
     def __init__(self, path=""):
         # load the optimized model
         print('creating base model...')
+        os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
         
         print('creating base model...')
         self.base_name = 'base40M-textvec'
@@ -49,6 +53,14 @@ class EndpointHandler():
         
         print('downloading upsampler checkpoint...')
         self.upsampler_model.load_state_dict(load_checkpoint('upsample', device))
+        
+        print('creating SDF model...')
+        self.sdf_name = 'sdf'
+        self.sdf_model = model_from_config(MODEL_CONFIGS[self.sdf_name], device)
+        self.sdf_model.eval()
+        
+        print('loading SDF model...')
+        self.sdf_model.load_state_dict(load_checkpoint(self.sdf_name, device))
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
@@ -113,6 +125,20 @@ class EndpointHandler():
         pc = sampler.output_to_point_clouds(samples)[0]
         print('type of pc: ', type(pc))
 
+        # Produce a mesh (with vertex colors)
+        mesh = marching_cubes_mesh(
+            pc=pc,
+            model=self.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)
+            print(mesh)
+
         pc_dict = {}
         
         data_list = pc.coords.tolist()
@@ -125,5 +151,6 @@ class EndpointHandler():
         # Serialize the dictionary to a JSON-formatted string
         channel_data = json.dumps(serializable_channels)
         pc_dict['channels'] = channel_data
-        
-        return pc_dict
+
+        #return pc_dict
+        return mesh