Update app.py

app.py

@@ -5,10 +5,35 @@ import gradio as gr
 import requests
 import tempfile
 
-device = torch.device("cpu")
-model = AutoModelForCausalLM.from_pretrained("microsoft/maira-2", trust_remote_code=True)
-processor = AutoProcessor.from_pretrained("microsoft/maira-2", trust_remote_code=True)
-model = model.eval().to(device)
+MODEL_STATE = {
+    "model": None,
+    "processor": None,
+    "authenticated": False
+}
+
+def login(hf_token):
+    """Authenticate and load the model"""
+    try:
+        MODEL_STATE.update({"model": None, "processor": None, "authenticated": False})
+        
+        MODEL_STATE["model"] = AutoModelForCausalLM.from_pretrained(
+            "microsoft/maira-2",
+            trust_remote_code=True,
+            use_auth_token=hf_token
+        )
+        MODEL_STATE["processor"] = AutoProcessor.from_pretrained(
+            "microsoft/maira-2",
+            trust_remote_code=True,
+            use_auth_token=hf_token
+        )
+        
+        MODEL_STATE["model"] = MODEL_STATE["model"].eval().to("cpu")
+        MODEL_STATE["authenticated"] = True
+        
+        return "🔓 Login successful! You can now use the model."
+    except Exception as e:
+        MODEL_STATE.update({"model": None, "processor": None, "authenticated": False})
+        return f"❌ Login failed: {str(e)}"
 
 def get_sample_data():
     """Download sample medical images and data"""
@@ -35,7 +60,6 @@ def save_temp_image(img):
     return temp_file.name
 
 def load_sample_findings():
-    """Load sample data for findings generation"""
     sample = get_sample_data()
     return [
         save_temp_image(sample["frontal"]),
@@ -47,22 +71,22 @@ def load_sample_findings():
     ]
 
 def load_sample_phrase():
-    """Load sample data for phrase grounding"""
     sample = get_sample_data()
     return [save_temp_image(sample["frontal"]), sample["phrase"]]
 
-def generate_report(frontal_path, lateral_path, indication, technique, comparison, 
+def generate_report(frontal_path, lateral_path, indication, technique, comparison,
                    prior_frontal_path, prior_lateral_path, prior_report, grounding):
-    """Generate radiology report with optional grounding"""
+    """Generate radiology report with authentication check"""
+    if not MODEL_STATE["authenticated"]:
+        return "⚠️ Please authenticate with your Hugging Face token first!"
+    
     try:
-        # Load images
         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
 
-        # Process inputs
-        processed = processor.format_and_preprocess_reporting_input(
+        processed = MODEL_STATE["processor"].format_and_preprocess_reporting_input(
             current_frontal=current_frontal,
             current_lateral=current_lateral,
             prior_frontal=prior_frontal,
@@ -73,88 +97,123 @@ def generate_report(frontal_path, lateral_path, indication, technique, compariso
             prior_report=prior_report or None,
             return_tensors="pt",
             get_grounding=grounding
-        ).to(device)
+        ).to("cpu")
 
-        # Generate report
-        outputs = model.generate(**processed, 
-                               max_new_tokens=450 if grounding else 300,
-                               use_cache=True)
+        outputs = MODEL_STATE["model"].generate(
+            **processed,
+            max_new_tokens=450 if grounding else 300,
+            use_cache=True
+        )
         
-        # Decode and format
         prompt_length = processed["input_ids"].shape[-1]
-        decoded = processor.decode(outputs[0][prompt_length:], skip_special_tokens=True)
-        return processor.convert_output_to_plaintext_or_grounded_sequence(decoded.lstrip())
+        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())
     
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"❌ Generation error: {str(e)}"
 
 def ground_phrase(frontal_path, phrase):
-    """Perform phrase grounding on image"""
+    """Perform phrase grounding with authentication check"""
+    if not MODEL_STATE["authenticated"]:
+        return "⚠️ Please authenticate with your Hugging Face token first!"
+    
     try:
         frontal = Image.open(frontal_path)
-        processed = processor.format_and_preprocess_phrase_grounding_input(
+        processed = MODEL_STATE["processor"].format_and_preprocess_phrase_grounding_input(
             frontal_image=frontal,
             phrase=phrase,
             return_tensors="pt"
-        ).to(device)
+        ).to("cpu")
         
-        outputs = model.generate(**processed, max_new_tokens=150, use_cache=True)
+        outputs = MODEL_STATE["model"].generate(
+            **processed,
+            max_new_tokens=150,
+            use_cache=True
+        )
         
         prompt_length = processed["input_ids"].shape[-1]
-        decoded = processor.decode(outputs[0][prompt_length:], skip_special_tokens=True)
-        return processor.convert_output_to_plaintext_or_grounded_sequence(decoded)
+        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)
     
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"❌ Grounding error: {str(e)}"
 
-# Gradio UI
-with gr.Blocks(title="MAIRA-2 Medical Imaging Assistant") as demo:
-    gr.Markdown("# MAIRA-2 Medical Imaging Assistant\nAI-powered radiology report generation and phrase grounding")
+with gr.Blocks(title="MAIRA-2 Medical Assistant") as demo:
+    gr.Markdown("""# MAIRA-2 Medical Assistant
+    **Authentication required** - You need a Hugging Face account and access token to use this model.
+    1. Get your access token from [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)
+    2. Request model access at [https://huggingface.co/microsoft/maira-2](https://huggingface.co/microsoft/maira-2)
+    3. Paste your token below to begin
+    """)
     
-    with gr.Tab("Report Generation"):
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown("## Current Study")
-                frontal = gr.Image(label="Frontal View", type="filepath")
-                lateral = gr.Image(label="Lateral View", type="filepath")
-                indication = gr.Textbox(label="Clinical Indication")
-                technique = gr.Textbox(label="Imaging Technique")
-                comparison = gr.Textbox(label="Comparison")
-                
-                gr.Markdown("## Prior Study (Optional)")
-                prior_frontal = gr.Image(label="Prior Frontal View", type="filepath")
-                prior_lateral = gr.Image(label="Prior Lateral View", type="filepath")
-                prior_report = gr.Textbox(label="Prior Report")
-                
-                grounding = gr.Checkbox(label="Include Grounding")
-                sample_btn = gr.Button("Load Sample Data")
-                
-            with gr.Column():
-                report_output = gr.Textbox(label="Generated Report", lines=10)
-                generate_btn = gr.Button("Generate Report")
-        
-        sample_btn.click(load_sample_findings,
-                        outputs=[frontal, lateral, indication, technique, comparison,
-                               prior_frontal, prior_lateral, prior_report, grounding])
-        generate_btn.click(generate_report,
-                         inputs=[frontal, lateral, indication, technique, comparison,
-                               prior_frontal, prior_lateral, prior_report, grounding],
-                         outputs=report_output)
+    with gr.Row():
+        hf_token = gr.Textbox(
+            label="Hugging Face Token",
+            placeholder="hf_xxxxxxxxxxxxxxxxxxxx",
+            type="password"
+        )
+        login_btn = gr.Button("Authenticate")
+        login_status = gr.Textbox(label="Authentication Status", interactive=False)
+    
+    login_btn.click(
+        login,
+        inputs=hf_token,
+        outputs=login_status
+    )
     
-    with gr.Tab("Phrase Grounding"):
-        with gr.Row():
-            with gr.Column():
-                pg_frontal = gr.Image(label="Frontal View", type="filepath")
-                phrase = gr.Textbox(label="Phrase to Ground")
-                pg_sample_btn = gr.Button("Load Sample Data")
-            with gr.Column():
-                pg_output = gr.Textbox(label="Grounding Result", lines=3)
-                pg_btn = gr.Button("Find Phrase")
+    with gr.Tabs():
+        with gr.Tab("Report Generation"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("## Current Study")
+                    frontal = gr.Image(label="Frontal View", type="filepath")
+                    lateral = gr.Image(label="Lateral View", type="filepath")
+                    indication = gr.Textbox(label="Clinical Indication")
+                    technique = gr.Textbox(label="Imaging Technique")
+                    comparison = gr.Textbox(label="Comparison")
+                    
+                    gr.Markdown("## Prior Study (Optional)")
+                    prior_frontal = gr.Image(label="Prior Frontal View", type="filepath")
+                    prior_lateral = gr.Image(label="Prior Lateral View", type="filepath")
+                    prior_report = gr.Textbox(label="Prior Report")
+                    
+                    grounding = gr.Checkbox(label="Include Grounding")
+                    sample_btn = gr.Button("Load Sample Data")
+                
+                with gr.Column():
+                    report_output = gr.Textbox(label="Generated Report", lines=10)
+                    generate_btn = gr.Button("Generate Report")
+            
+            sample_btn.click(
+                load_sample_findings,
+                outputs=[frontal, lateral, indication, technique, comparison,
+                       prior_frontal, prior_lateral, prior_report, grounding]
+            )
+            generate_btn.click(
+                generate_report,
+                inputs=[frontal, lateral, indication, technique, comparison,
+                      prior_frontal, prior_lateral, prior_report, grounding],
+                outputs=report_output
+            )
         
-        pg_sample_btn.click(load_sample_phrase,
-                           outputs=[pg_frontal, phrase])
-        pg_btn.click(ground_phrase,
-                    inputs=[pg_frontal, phrase],
-                    outputs=pg_output)
+        with gr.Tab("Phrase Grounding"):
+            with gr.Row():
+                with gr.Column():
+                    pg_frontal = gr.Image(label="Frontal View", type="filepath")
+                    phrase = gr.Textbox(label="Phrase to Ground")
+                    pg_sample_btn = gr.Button("Load Sample Data")
+                with gr.Column():
+                    pg_output = gr.Textbox(label="Grounding Result", lines=3)
+                    pg_btn = gr.Button("Find Phrase")
+            
+            pg_sample_btn.click(
+                load_sample_phrase,
+                outputs=[pg_frontal, phrase]
+            )
+            pg_btn.click(
+                ground_phrase,
+                inputs=[pg_frontal, phrase],
+                outputs=pg_output
+            )
 
 demo.launch()