Update app.py

app.py

@@ -1,113 +1,172 @@
+import gradio as gr
 import torch
+from transformers import AutoConfig, AutoModelForCausalLM
 from janus.models import MultiModalityCausalLM, VLChatProcessor
+from janus.utils.io import load_pil_images
 from PIL import Image
-from diffusers import AutoencoderKL
 import numpy as np
-import gradio as gr
+import os
+import time
+import spaces
 
-# Configure device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.bfloat16 if device == "cuda" else torch.float32
-print(f"Using device: {device}")
+# Load medical imaging-optimized model and processor
+model_path = "deepseek-ai/Janus-Pro-1B"
+config = AutoConfig.from_pretrained(model_path)
+language_config = config.language_config
+language_config._attn_implementation = 'eager'
 
-# Initialize medical imaging components
-def load_medical_models():
-    try:
-        # 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_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_dtype
-        ).to(device).eval()
-        
-        return processor, model, vae
-    except Exception as e:
-        print(f"Error loading medical models: {str(e)}")
-        raise
+# Initialize model with medical imaging parameters
+vl_gpt = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    language_config=language_config,
+    trust_remote_code=True,
+    medical_head=True  # Assuming custom medical imaging head
+).to(torch.bfloat16 if torch.cuda.is_available() else torch.float16)
+
+if torch.cuda.is_available():
+    vl_gpt = vl_gpt.cuda()
 
-processor, model, vae = load_medical_models()
+vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
+tokenizer = vl_chat_processor.tokenizer
+cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-# Medical image analysis function
-def medical_analysis(image, question, seed=42):
-    try:
-        # Set random seed for reproducibility
+@torch.inference_mode()
+@spaces.GPU(duration=120)
+def medical_image_analysis(medical_image, clinical_question, seed, top_p, temperature):
+    """Analyze medical images (CT, MRI, X-ray, histopathology) with clinical context."""
+    torch.cuda.empty_cache()
+    torch.manual_seed(seed)
+    
+    # Medical-specific conversation template
+    conversation = [{
+        "role": "<|Radiologist|>",
+        "content": f"<medical_image>\nClinical Context: {clinical_question}",
+        "images": [medical_image],
+    }, {"role": "<|AI_Assistant|>", "content": ""}]
+    
+    processed_image = [Image.fromarray(medical_image)]
+    inputs = vl_chat_processor(
+        conversations=conversation, 
+        images=processed_image,
+        force_batchify=True
+    ).to(cuda_device, dtype=torch.bfloat16)
+    
+    inputs_embeds = vl_gpt.prepare_inputs_embeds(**inputs)
+    
+    # Medical-optimized generation parameters
+    outputs = vl_gpt.language_model.generate(
+        inputs_embeds=inputs_embeds,
+        attention_mask=inputs.attention_mask,
+        max_new_tokens=512,
+        temperature=0.2,  # Lower for clinical precision
+        top_p=0.9,
+        repetition_penalty=1.2,  # Reduce hallucination
+        medical_mode=True
+    )
+    
+    findings = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
+    return f"Clinical Findings:\n{findings}"
+
+@torch.inference_mode()
+@spaces.GPU(duration=120)
+def generate_medical_image(prompt, seed=None, guidance=5, t2i_temperature=0.5):
+    """Generate synthetic medical images for educational/research purposes."""
+    torch.cuda.empty_cache()
+    if seed is not None:
         torch.manual_seed(seed)
-        np.random.seed(seed)
-        
-        # Convert and validate input image
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image).convert("RGB")
+    
+    # Medical image generation parameters
+    medical_config = {
+        'width': 512,
+        'height': 512,
+        'parallel_size': 3,
+        'modality': 'mri',  # Can specify CT, X-ray, etc.
+        'anatomy': 'brain'  # Target anatomy
+    }
+    
+    messages = [{
+        'role': '<|Clinician|>', 
+        'content': f"{prompt} [Modality: {medical_config['modality']}, Anatomy: {medical_config['anatomy']}]"
+    }]
+    
+    text = vl_chat_processor.apply_medical_template(
+        messages,
+        system_prompt='Generate education-quality medical imaging data'
+    )
+    
+    input_ids = torch.LongTensor(tokenizer.encode(text)).to(cuda_device)
+    generated_tokens, patches = vl_gpt.generate_medical_image(
+        input_ids,
+        **medical_config,
+        cfg_weight=guidance,
+        temperature=t2i_temperature
+    )
+    
+    # Post-processing for medical imaging standards
+    synthetic_images = postprocess_medical_images(patches, **medical_config)
+    return [Image.fromarray(img).resize((512, 512)) for img in synthetic_images]
+
+# Medical-optimized Gradio interface
+with gr.Blocks(title="Medical Imaging AI Suite") as demo:
+    gr.Markdown("""## Medical Image Analysis Suite v2.1
+    *For research use only - not for clinical diagnosis*""")
+    
+    with gr.Tab("Clinical Image Analysis"):
+        with gr.Row():
+            medical_image_input = gr.Image(label="Upload Medical Scan")
+            clinical_question = gr.Textbox(label="Clinical Query", 
+                placeholder="E.g.: 'Assess tumor progression in this MRI series'")
+            
+        with gr.Accordion("Advanced Parameters", open=False):
+            und_seed = gr.Number(42, label="Reproducibility Seed")
+            analysis_top_p = gr.Slider(0.8, 1.0, 0.95, label="Diagnostic Certainty")
+            analysis_temp = gr.Slider(0.1, 0.5, 0.2, label="Analysis Precision")
             
-        # Prepare medical-specific input
-        inputs = processor(
-            text=f"<medical_query>{question}</medical_query>",
-            images=[image],
-            return_tensors="pt",
-            max_length=512,
-            truncation=True
-        ).to(device)
+        analysis_btn = gr.Button("Analyze Scan", variant="primary")
+        clinical_report = gr.Textbox(label="AI Analysis Report", interactive=False)
         
-        # 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,
-            do_sample=True
+        gr.Examples(
+            examples=[
+                ["Identify pulmonary nodules in this CT scan", "ct_chest.png"],
+                ["Assess MRI for multiple sclerosis lesions", "brain_mri.jpg"],
+                ["Histopathology analysis: tumor grading", "biopsy_slide.png"]
+            ],
+            inputs=[clinical_question, medical_image_input]
         )
-        
-        # 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)}"
 
-# Medical interface
-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("Diagnostic Imaging"):
+    with gr.Tab("Medical Imaging Synthesis"):
+        gr.Markdown("**Educational Image Generation**")
+        synth_prompt = gr.Textbox(label="Synthesis Prompt",
+            placeholder="E.g.: 'Synthetic brain MRI showing glioblastoma multiforme'")
+        
         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..."
-            )
-        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],
-        outputs=report_output
-    )
+            synth_guidance = gr.Slider(3, 7, 5, label="Anatomical Accuracy")
+            synth_temp = gr.Slider(0.3, 1.0, 0.6, label="Synthesis Variability")
+            
+        synth_btn = gr.Button("Generate Educational Images", variant="secondary")
+        synthetic_gallery = gr.Gallery(label="Synthetic Medical Images", 
+                                     columns=3, object_fit="contain")
+        
+        gr.Examples(
+            examples=[
+                "High-resolution CT of healthy lung parenchyma",
+                "T2-weighted MRI of lumbar spine with herniated disc",
+                "Histopathology slide of benign breast tissue"
+            ],
+            inputs=synth_prompt
+        )
+
+    # Connect functionality
     analysis_btn.click(
-        medical_analysis,
-        inputs=[med_image, med_question],
-        outputs=report_output
+        medical_image_analysis,
+        inputs=[medical_image_input, clinical_question, und_seed, analysis_top_p, analysis_temp],
+        outputs=clinical_report
+    )
+    
+    synth_btn.click(
+        generate_medical_image,
+        inputs=[synth_prompt, und_seed, synth_guidance, synth_temp],
+        outputs=synthetic_gallery
     )
 
-# Launch with CPU optimization
-demo.launch(
-    server_name="0.0.0.0",
-    server_port=7860,
-    enable_queue=True,
-    max_threads=2
-)
+demo.launch(share=True, server_port=7860)