Update app.py

app.py

@@ -116,7 +116,7 @@ def server(input: Inputs, output: Outputs, session: Session):
             return None  # Return an empty string if no image is uploaded
 
     @reactive.Calc
-    def generate_input_points():
+    def generate_input_points(array_size = 256):
         """
         input_points (torch.FloatTensor of shape (batch_size, num_points, 2)) —
         Input 2D spatial points, this is used by the prompt encoder to encode the prompt.
@@ -132,8 +132,8 @@ def server(input: Inputs, output: Outputs, session: Session):
         embedding will be skipped for these points using the labels.
         
         """
-        # Define the size of your array
-        array_size = 256
+        # # Define the size of your array
+        # array_size = 256
         
         # Define the size of your grid
         grid_size = 10
@@ -177,7 +177,8 @@ def server(input: Inputs, output: Outputs, session: Session):
 
         """ Prepare Inputs """
         # get input points prompt (grid of points)
-        input_points = generate_input_points(image)
+        array_size = max(image.shape[0], image.shape[1])
+        input_points = generate_input_points(array_size)
         
         # prepare image and prompt for the model
         inputs = processor(image, input_points=input_points, return_tensors="pt")