Update app.py

app.py

@@ -211,9 +211,7 @@ def count_objects(image, pos_caption, neg_caption, box_threshold, point_radius,
     
 
     if use_neg:
-        print('neg_caption: ', neg_caption)
-        if neg_caption is None:
-            neg_caption = "None."
+        # print('neg_caption: ', neg_caption)
         neg_inputs = processor(
             images=image,
             text=neg_caption,
@@ -231,6 +229,22 @@ def count_objects(image, pos_caption, neg_caption, box_threshold, point_radius,
         pos_inputs['neg_input_ids'] = neg_inputs['input_ids']
         pos_inputs['use_neg'] = True
     else:
+        neg_caption = "None."
+        neg_inputs = processor(
+            images=image,
+            text=neg_caption,
+            return_tensors="pt",
+            padding=True
+        )
+        neg_inputs = {k: v.to(device) for k, v in neg_inputs.items()}
+        neg_inputs['pixel_values'] = neg_inputs['pixel_values'].to(torch.bfloat16)
+
+        # Add negative inputs to positive inputs dict
+        pos_inputs['neg_token_type_ids'] = neg_inputs['token_type_ids']
+        pos_inputs['neg_attention_mask'] = neg_inputs['attention_mask']
+        pos_inputs['neg_pixel_mask'] = neg_inputs['pixel_mask']
+        pos_inputs['neg_pixel_values'] = neg_inputs['pixel_values']
+        pos_inputs['neg_input_ids'] = neg_inputs['input_ids']
         pos_inputs['use_neg'] = False
 
     # Run inference