Update handler.py

handler.py

@@ -23,7 +23,10 @@ class EndpointHandler():
     def __init__(self, path=""):
         # load the optimized model
         print('creating base model...')
-        self.base_name = 'base40M-textvec'
+        
+        #self.base_name = 'base40M-textvec'
+        self.base_name = 'base40M'
+        
         self.base_model = model_from_config(MODEL_CONFIGS[self.base_name], device)
         self.base_model.eval()
         self.base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[self.base_name])
@@ -47,6 +50,15 @@ class EndpointHandler():
         Return:
             A :obj:`dict`:. plotly json Data
         """
+
+        #Checks if an image key has been provided, and if so, uses the image data instead of text input
+        if "image" in data:
+            image_data = data.pop("image")
+            use_image = True
+            print('image data found')
+        else
+            print('no image data found')
+        
         inputs = data.pop("inputs", data)
 
         sampler = PointCloudSampler(