Update app.py

app.py

@@ -5,23 +5,23 @@ from diffusers import AutoencoderKL
 import numpy as np
 import gradio as gr
 
-# Configure device
+# Configure device and attention implementation
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
+attn_implementation = "flash_attention_2" if device == "cuda" else "eager"
+print(f"Using device: {device} with {attn_implementation}")
 
 # Initialize medical imaging components
 def load_medical_models():
     try:
-        # Load processor and tokenizer
         processor = VLChatProcessor.from_pretrained("deepseek-ai/Janus-1.3B")
         
-        # Load base model
         model = MultiModalityCausalLM.from_pretrained(
             "deepseek-ai/Janus-1.3B",
-            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
+            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
+            attn_implementation=attn_implementation,
+            use_flash_attention_2=(attn_implementation == "flash_attention_2")
         ).to(device).eval()
         
-        # Load VAE for image processing
         vae = AutoencoderKL.from_pretrained(
             "stabilityai/sdxl-vae",
             torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
@@ -29,62 +29,61 @@ def load_medical_models():
         
         return processor, model, vae
     except Exception as e:
-        print(f"Error loading models: {str(e)}")
+        print(f"Error loading medical models: {str(e)}")
         raise
 
 processor, model, vae = load_medical_models()
 
-# Medical image analysis function
-def medical_analysis(image, question, seed=42, top_p=0.95, temperature=0.1):
+# Medical image analysis function with attention control
+def medical_analysis(image, question, seed=42):
     try:
-        # Set random seed for reproducibility
         torch.manual_seed(seed)
         np.random.seed(seed)
         
-        # Prepare inputs
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image).convert("RGB")
             
         inputs = processor(
-            text=question,
+            text=f"<medical_query>{question}</medical_query>",
             images=[image],
             return_tensors="pt"
         ).to(device)
         
-        # Generate analysis
         outputs = model.generate(
             inputs.input_ids,
             attention_mask=inputs.attention_mask,
             max_new_tokens=512,
-            temperature=temperature,
-            top_p=top_p
+            temperature=0.1,
+            top_p=0.95,
+            pad_token_id=processor.tokenizer.eos_token_id
         )
         
         return processor.decode(outputs[0], skip_special_tokens=True)
     except Exception as e:
-        return f"Analysis error: {str(e)}"
+        return f"Radiology analysis error: {str(e)}"
 
 # Medical interface
-with gr.Blocks(title="Medical Imaging Assistant") as demo:
-    gr.Markdown("# Medical Imaging AI Assistant")
+with gr.Blocks(title="Medical Imaging Assistant", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""# AI Radiology Assistant
+                **CT/MRI/X-ray Analysis System**""")
     
-    with gr.Tab("Analysis"):
+    with gr.Tab("Diagnostic Imaging"):
         with gr.Row():
-            med_image = gr.Image(label="Input Image", type="pil")
-            med_question = gr.Textbox(label="Clinical Query")
-        analysis_output = gr.Textbox(label="Findings")
-        gr.Examples(
-            examples=[
-                ["ultrasound_sample.jpg", "Identify any abnormalities in this ultrasound"],
-                ["xray_sample.jpg", "Describe the bone structure visible in this X-ray"]
-            ],
-            inputs=[med_image, med_question]
-        )
+            med_image = gr.Image(label="DICOM Image", type="pil")
+            med_question = gr.Textbox(label="Clinical Query", 
+                                    placeholder="Describe findings in this CT scan...")
+        analysis_btn = gr.Button("Analyze", variant="primary")
+        report_output = gr.Textbox(label="Radiology Report", interactive=False)
     
     med_question.submit(
         medical_analysis,
         inputs=[med_image, med_question],
-        outputs=analysis_output
+        outputs=report_output
+    )
+    analysis_btn.click(
+        medical_analysis,
+        inputs=[med_image, med_question],
+        outputs=report_output
     )
 
-demo.launch()
+demo.launch(server_name="0.0.0.0", server_port=7860)