Update app.py

app.py

@@ -80,31 +80,30 @@ def generate_report(frontal_path, lateral_path, indication, technique, compariso
     if not MODEL_STATE["authenticated"]:
         return "⚠️ Please authenticate with your Hugging Face token first!"
     
+    if not frontal_path or not lateral_path:
+        return "❌ Please upload both the frontal and lateral images for the current study."
+    
     try:
-        current_frontal = Image.open(frontal_path) if frontal_path else None
-        current_lateral = Image.open(lateral_path) if lateral_path else None
+        current_frontal = Image.open(frontal_path)
+        current_lateral = Image.open(lateral_path)
+        prior_frontal = Image.open(prior_frontal_path) if prior_frontal_path else None
+        prior_lateral = Image.open(prior_lateral_path) if prior_lateral_path else None
+
+        processed = MODEL_STATE["processor"].format_and_preprocess_reporting_input(
+            current_frontal=current_frontal,
+            current_lateral=current_lateral,
+            prior_frontal=prior_frontal,
+            prior_lateral=prior_lateral,
+            indication=indication,
+            technique=technique,
+            comparison=comparison,
+            prior_report=prior_report or None,
+            return_tensors="pt",
+            get_grounding=grounding
+        ).to("cpu")
+
+        processed.pop("image_sizes", None)
 
-        kwargs = {
-            "current_frontal": current_frontal,
-            "current_lateral": current_lateral,
-            "indication": indication,
-            "technique": technique,
-            "comparison": comparison,
-            "prior_report": prior_report or None,
-            "return_tensors": "pt",
-            "get_grounding": grounding
-        }
-        if prior_frontal_path:
-            kwargs["prior_frontal"] = Image.open(prior_frontal_path)
-        if prior_lateral_path:
-            kwargs["prior_lateral"] = Image.open(prior_lateral_path)
-        
-        processed = MODEL_STATE["processor"].format_and_preprocess_reporting_input(**kwargs).to("cpu")
-        
-        # Remove any unsupported keys (e.g. image_sizes) if present
-        if "image_sizes" in processed:
-            processed.pop("image_sizes")
-        
         outputs = MODEL_STATE["model"].generate(
             **processed,
             max_new_tokens=450 if grounding else 300,
@@ -123,17 +122,20 @@ def ground_phrase(frontal_path, phrase):
     if not MODEL_STATE["authenticated"]:
         return "⚠️ Please authenticate with your Hugging Face token first!"
     
+    # Check that the required image is provided.
+    if not frontal_path:
+        return "❌ Please upload the frontal image for phrase grounding."
+    
     try:
-        frontal = Image.open(frontal_path) if frontal_path else None
+        frontal = Image.open(frontal_path)
         processed = MODEL_STATE["processor"].format_and_preprocess_phrase_grounding_input(
             frontal_image=frontal,
             phrase=phrase,
             return_tensors="pt"
         ).to("cpu")
         
-        # Remove unsupported keyword if present
-        if "image_sizes" in processed:
-            processed.pop("image_sizes")
+        # Remove the unexpected key if present.
+        processed.pop("image_sizes", None)
         
         outputs = MODEL_STATE["model"].generate(
             **processed,