pass max_new_tokens

app.py

@@ -15,8 +15,8 @@ COLORS = ['#4285f4', '#db4437', '#f4b400', '#0f9d58', '#e48ef1']
 ##### Parse segmentation output tokens into masks
 ##### Also returns bounding boxes with their labels
 
-def parse_segmentation(input_image, input_text):
-  out = pali_gemma_model.infer(input_image, input_text, max_new_tokens=100)
+def parse_segmentation(input_image, input_text, max_new_tokens=100):
+  out = pali_gemma_model.infer(input_image, input_text, max_new_tokens=max_new_tokens)
   objs = extract_objs(out.lstrip("\n"), input_image.size[0], input_image.size[1], unique_labels=True)
   labels = set(obj.get('name') for obj in objs if obj.get('name'))
   color_map = {l: COLORS[i % len(COLORS)] for i, l in enumerate(labels)}
@@ -72,7 +72,7 @@ with gr.Blocks(css="style.css") as demo:
     ]
     seg_btn.click(
         fn=parse_segmentation,
-        inputs=seg_inputs,
+        inputs=[seg_inputs, "max_new_tokens"],
         outputs=seg_outputs,
     )
   with gr.Tab("Text Generation"):