Update app.py

app.py

@@ -105,9 +105,10 @@ def generate_report(frontal_path, lateral_path, indication, technique, compariso
             get_grounding=grounding
         ).to("cpu")
         
-        
-        if "image_sizes" in processed:
-            processed.pop("image_sizes")
+        # Remove all image size related keys
+        image_size_keys = [k for k in processed.keys() if "image_sizes" in k]
+        for k in image_size_keys:
+            processed.pop(k)
         
         outputs = MODEL_STATE["model"].generate(
             **processed,
@@ -138,9 +139,10 @@ def ground_phrase(frontal_path, phrase):
             return_tensors="pt"
         ).to("cpu")
         
-        
-        if "image_sizes" in processed:
-            processed.pop("image_sizes")
+        # Remove all image size related keys
+        image_size_keys = [k for k in processed.keys() if "image_sizes" in k]
+        for k in image_size_keys:
+            processed.pop(k)
         
         outputs = MODEL_STATE["model"].generate(
             **processed,