Update app.py

app.py

@@ -104,16 +104,18 @@ def generate_report(frontal_path, lateral_path, indication, technique, compariso
             return_tensors="pt",
             get_grounding=grounding
         ).to("cpu")
-
-        processed_inputs = {k: v for k, v in processed.items() if k != 'image_sizes'}
-
+        
+        
+        if "image_sizes" in processed:
+            processed.pop("image_sizes")
+        
         outputs = MODEL_STATE["model"].generate(
-            **processed_inputs,
+            **processed,
             max_new_tokens=450 if grounding else 300,
             use_cache=True
         )
         
-        prompt_length = processed_inputs["input_ids"].shape[-1]
+        prompt_length = processed["input_ids"].shape[-1]
         decoded = MODEL_STATE["processor"].decode(outputs[0][prompt_length:], skip_special_tokens=True)
         return MODEL_STATE["processor"].convert_output_to_plaintext_or_grounded_sequence(decoded.lstrip())
     
@@ -135,16 +137,18 @@ def ground_phrase(frontal_path, phrase):
             phrase=phrase,
             return_tensors="pt"
         ).to("cpu")
-
-        processed_inputs = {k: v for k, v in processed.items() if k != 'image_sizes'}
+        
+        
+        if "image_sizes" in processed:
+            processed.pop("image_sizes")
         
         outputs = MODEL_STATE["model"].generate(
-            **processed_inputs,
+            **processed,
             max_new_tokens=150,
             use_cache=True
         )
         
-        prompt_length = processed_inputs["input_ids"].shape[-1]
+        prompt_length = processed["input_ids"].shape[-1]
         decoded = MODEL_STATE["processor"].decode(outputs[0][prompt_length:], skip_special_tokens=True)
         return MODEL_STATE["processor"].convert_output_to_plaintext_or_grounded_sequence(decoded)
     
@@ -229,4 +233,4 @@ with gr.Blocks(title="MAIRA-2 Medical Assistant") as demo:
                 outputs=pg_output
             )
 
-demo.launch()
+demo.launch()