Updated Handler with support for custom 'commands' and seperated code into methods for more efficient calling.

handler.py

@@ -62,15 +62,36 @@ class EndpointHandler():
         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]]]:
-        """
-        Args:
-            data (:obj:):
-                includes the input data and the parameters for the inference.
-        Return:
-            A :obj:`dict`:. plotly json Data
-        """
+    def __call__(self, input_data: Any) -> Any:
+
+        command = "null"
+
+        if "command" in input_data:
+            command = input_data["command"]
+
+        print(f"the command is: {command}")
+
+        #Assume the user app is still running the old version, and send the data back as it is being expected
+        #Currently, the App expects a .ply Mesh to be sent back, and will not have a command input sent with it
+        if command == "null":
+            temp_pc = self.generate_point_cloud(input_data)
+            return self.generate_mesh_from_pc(temp_pc)
+        elif command == "generate_pc":
+            return self.generate_point_cloud(input_data)  
+        elif command == "generate_mesh":
+            return self.generate_mesh_from_pc(input_data["raw_pc"])
+        elif command == "status":
+            return self.check_status()
+            
+        
+
+
+    def check_status(self) -> bool:
+        return self.active
+
 
+    def generate_point_cloud(self, data: Any) -> Dict[str, Dict[str, float]]:
+        print("generate pc called...")
         use_image = False
         
         #Checks if an image key has been provided, and if so, uses the image data instead of text input
@@ -123,11 +144,40 @@ class EndpointHandler():
         #image = self.pipe(inputs, guidance_scale=7.5)["sample"][0]  
 
         pc = sampler.output_to_point_clouds(samples)[0]
-        print('type of pc: ', type(pc))
 
-        # Produce a mesh (with vertex colors)
+        pc_dict = {}
+
+        data_list = pc.coords.tolist()
+        json_string = json.dumps(data_list)
+        pc_dict['data'] = json_string
+        
+        # Convert NumPy arrays to Python lists for serializing
+        serializable_channels = {key: value.tolist() for key, value in pc.channels.items()}
+        
+        # Serialize the dictionary to a JSON-formatted string
+        channel_data = json.dumps(serializable_channels)
+        pc_dict['channels'] = channel_data
+        
+        return pc_dict
+
+
+    def generate_mesh_from_pc(self, pc_data: Any) -> Any:
+         # Produce a mesh (with vertex colors)
+        print("generate mesh called...")
+
+        #De-serialize both the coords and channel data
+        coords_list = json.loads(pc_data['data'])
+        channels_dict = json.loads(pc_data['channels'])
+
+        # Reconstruct the PointCloud object
+        # Make sure to use .items() for the dictionary to output the key-value pairs together
+        point_cloud = PointCloud(
+            coords=np.array(coords_list, dtype=np.float32),
+            channels={name: np.array(array, dtype=np.float32) for name, array in channels_dict.items()}
+        )
+        
         mesh = marching_cubes_mesh(
-            pc=pc,
+            pc=point_cloud,
             model=self.sdf_model,
             batch_size=4096,
             grid_size=32, # increase to 128 for resolution used in evals
@@ -138,19 +188,6 @@ class EndpointHandler():
         with open('mesh.ply', 'wb') as f:
             mesh.write_ply(f)
             print(mesh)
-
-        pc_dict = {}
-        
-        data_list = pc.coords.tolist()
-        json_string = json.dumps(data_list)
-        pc_dict['data'] = json_string
-
-        # Convert NumPy arrays to Python lists for serializing
-        serializable_channels = {key: value.tolist() for key, value in pc.channels.items()}
-
-        # Serialize the dictionary to a JSON-formatted string
-        channel_data = json.dumps(serializable_channels)
-        pc_dict['channels'] = channel_data
-
-        #return pc_dict
+            
         return mesh
+