major change, used biomed_llama 7b

app.py

@@ -17,6 +17,7 @@ import sys
 from pathlib import Path
 from huggingface_hub import login
 # from dotenv import load_dotenv
+from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 
 # For Hugging Face Spaces, secrets are automatically loaded as environment variables
 token = os.getenv("HF_TOKEN")
@@ -188,8 +189,32 @@ def initialize_model():
         )
     return model
 
+def initialize_llm():
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_compute_dtype=torch.float16
+    )
+    
+    model = AutoModel.from_pretrained(
+        "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
+        quantization_config=bnb_config,
+        device_map="auto",
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+        attn_implementation="flash_attention_2"
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained(
+        "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
+        trust_remote_code=True
+    )
+    
+    return model, tokenizer
 
 model = initialize_model()
+llm_model, llm_tokenizer = initialize_llm()
 
 def update_example_prompts(modality):
     if modality in MODALITY_PROMPTS:
@@ -211,39 +236,53 @@ def process_image(image_path, text_prompts, modality):
         if not modality:
             raise ValueError("Please select a modality")
 
+        # Original BiomedParse processing
         image = read_rgb(image_path)
         text_prompts = [prompt.strip() for prompt in text_prompts.split(',')]
-        
-        # Run inference
         pred_masks = interactive_infer_image(model, Image.fromarray(image), text_prompts)
         
         # Prepare outputs
         results = []
-        p_values = []
-
+        analysis_results = []
+        
+        # Process with BiomedParse
         for i, prompt in enumerate(text_prompts):
-            # Calculate p-value for the selected modality
-            print("PROMPT: ", prompt, flush=True)   
             p_value = check_mask_stats(image, pred_masks[i] * 255, modality, prompt)
-            p_values.append(f"P-value for '{prompt}' ({modality}): {p_value:.4f}")
-
-            # Overlay predictions on the image
+            analysis_results.append(f"P-value for '{prompt}' ({modality}): {p_value:.4f}")
+            
             overlay_image = image.copy()
-            overlay_image[pred_masks[i] > 0.5] = [255, 0, 0]  # Highlight predictions in red
+            overlay_image[pred_masks[i] > 0.5] = [255, 0, 0]
             results.append(overlay_image)
+        
+        # Process with LLM
+        pil_image = Image.fromarray(image)
+        question = 'Give the modality, organ, analysis, abnormalities (if any), treatment (if abnormalities are present)?'
+        msgs = [{'role': 'user', 'content': [pil_image, question]}]
+        
+        llm_response = ""
+        for new_text in llm_model.chat(
+            image=pil_image,
+            msgs=msgs,
+            tokenizer=llm_tokenizer,
+            sampling=True,
+            temperature=0.95,
+            stream=True
+        ):
+            llm_response += new_text
+            
+        # Combine both analyses
+        combined_analysis = "\n\n".join([
+            "BiomedParse Analysis:",
+            "\n".join(analysis_results),
+            "\nLLM Analysis:",
+            llm_response
+        ])
+        
+        return results, combined_analysis
 
-            return results, "\n".join(p_values)
-
-    except ValueError as ve:
-        # Handle validation errors
-        return None, f"⚠️ Input Error: {str(ve)}"
-    except torch.cuda.OutOfMemoryError:
-        # Handle CUDA out of memory errors
-        return None, "⚠️ Error: GPU memory exceeded. Please try with a smaller image."
     except Exception as e:
-        # Handle all other errors
-        error_msg = f"��️ An error occurred: {str(e)}"
-        print(f"Error details: {str(e)}", flush=True)  # For logging
+        error_msg = f"⚠️ An error occurred: {str(e)}"
+        print(f"Error details: {str(e)}", flush=True)
         return None, error_msg
 
 # Define Gradio interface