Update app.py

app.py

@@ -177,7 +177,7 @@ def server(input: Inputs, output: Outputs, session: Session):
 
         """ Prepare Inputs """
         # get input points prompt (grid of points)
-        array_size = max(image.shape[0], image.shape[1])
+        array_size = max(image.size)
         input_points = generate_input_points(array_size)
         
         # prepare image and prompt for the model