Update handler.py

handler.py

@@ -60,7 +60,7 @@ class EndpointHandler():
 
         #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")
+            image_data_encoded = data.pop("image")
             use_image = True
             print('image data found')
         else:
@@ -70,7 +70,7 @@ class EndpointHandler():
         inputs = data.pop("inputs", data)
 
        if use_image:
-           sampler = PointCloudSampler(
+            sampler = PointCloudSampler(
                 device=device,
                 models=[base_model, upsampler_model],
                 diffusions=[base_diffusion, upsampler_diffusion],
@@ -80,7 +80,10 @@ class EndpointHandler():
             )
 
             # Load an image to condition on.
-            img = Image.open('example_data/cube_stack.jpg')
+            image_data = base64.b64decode(image_data_encoded)
+    
+            # Convert bytes to PIL Image
+            img = Image.open(BytesIO(image_data))
        else:
             sampler = PointCloudSampler(
                 device=device,