Update app.py

app.py

@@ -20,14 +20,15 @@ PROCESSOR_IDS = {
     "Model 2 (WaveUI 896)": "google/paligemma-3b-pt-896"
 }
 
+# Device configuration
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Load models and processors
 models = {name: PaliGemmaForConditionalGeneration.from_pretrained(model_id).eval().to(device)
           for name, model_id in MODEL_IDS.items()}
 processors = {name: PaliGemmaProcessor.from_pretrained(processor_id)
               for name, processor_id in PROCESSOR_IDS.items()}
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 ###### Transformers Inference
 @spaces.GPU
 def infer(
@@ -62,7 +63,7 @@ def parse_segmentation(input_image, input_text, model_choice):
                 obj['name'] or '',
             )
             for obj in objs
-            if 'mask' in obj or 'xyxy' in obj
+            if 'mask' in obj or 'xyxy'
         ],
     )
     has_annotations = bool(annotated_img[1])
@@ -108,7 +109,6 @@ with gr.Blocks(css="style.css") as demo:
             outputs=seg_outputs,
         )
 
-
 _SEGMENT_DETECT_RE = re.compile(
     r'(.*?)' +
     r'<loc(\d{4})>' * 4 + r'\s*' +