Update app.py

app.py

@@ -5,26 +5,32 @@ from diffusers import AutoencoderKL
 import numpy as np
 import gradio as gr
 
-# Configure device and attention implementation
+# Configure device and disable FlashAttention
 device = "cuda" if torch.cuda.is_available() else "cpu"
-attn_implementation = "flash_attention_2" if device == "cuda" else "eager"
-print(f"Using device: {device} with {attn_implementation}")
+torch_dtype = torch.bfloat16 if device == "cuda" else torch.float32
+print(f"Using device: {device}")
 
 # Initialize medical imaging components
 def load_medical_models():
     try:
-        processor = VLChatProcessor.from_pretrained("deepseek-ai/Janus-1.3B")
+        # Load processor with medical-specific configuration
+        processor = VLChatProcessor.from_pretrained(
+            "deepseek-ai/Janus-1.3B",
+            medical_mode=True
+        )
         
+        # Load model with CPU/GPU optimization
         model = MultiModalityCausalLM.from_pretrained(
             "deepseek-ai/Janus-1.3B",
-            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
-            attn_implementation=attn_implementation,
-            use_flash_attention_2=(attn_implementation == "flash_attention_2")
+            torch_dtype=torch_dtype,
+            attn_implementation="eager",  # Force standard attention
+            low_cpu_mem_usage=True
         ).to(device).eval()
         
+        # Load VAE with reduced precision
         vae = AutoencoderKL.from_pretrained(
             "stabilityai/sdxl-vae",
-            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
+            torch_dtype=torch_dtype
         ).to(device).eval()
         
         return processor, model, vae
@@ -34,31 +40,40 @@ def load_medical_models():
 
 processor, model, vae = load_medical_models()
 
-# Medical image analysis function with attention control
+# Medical image analysis function
 def medical_analysis(image, question, seed=42):
     try:
+        # Set random seed for reproducibility
         torch.manual_seed(seed)
         np.random.seed(seed)
         
+        # Convert and validate input image
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image).convert("RGB")
             
+        # Prepare medical-specific input
         inputs = processor(
             text=f"<medical_query>{question}</medical_query>",
             images=[image],
-            return_tensors="pt"
+            return_tensors="pt",
+            max_length=512,
+            truncation=True
         ).to(device)
         
+        # Generate medical analysis
         outputs = model.generate(
             inputs.input_ids,
             attention_mask=inputs.attention_mask,
             max_new_tokens=512,
             temperature=0.1,
             top_p=0.95,
-            pad_token_id=processor.tokenizer.eos_token_id
+            pad_token_id=processor.tokenizer.eos_token_id,
+            do_sample=True
         )
         
-        return processor.decode(outputs[0], skip_special_tokens=True)
+        # Clean and return medical report
+        report = processor.decode(outputs[0], skip_special_tokens=True)
+        return report.replace("##MEDICAL_REPORT##", "").strip()
     except Exception as e:
         return f"Radiology analysis error: {str(e)}"
 
@@ -70,11 +85,14 @@ with gr.Blocks(title="Medical Imaging Assistant", theme=gr.themes.Soft()) as dem
     with gr.Tab("Diagnostic Imaging"):
         with gr.Row():
             med_image = gr.Image(label="DICOM Image", type="pil")
-            med_question = gr.Textbox(label="Clinical Query", 
-                                    placeholder="Describe findings in this CT scan...")
+            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)
     
+    # Connect components
     med_question.submit(
         medical_analysis,
         inputs=[med_image, med_question],
@@ -86,4 +104,10 @@ with gr.Blocks(title="Medical Imaging Assistant", theme=gr.themes.Soft()) as dem
         outputs=report_output
     )
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+# Launch with CPU optimization
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+    enable_queue=True,
+    max_threads=2
+)