major change, used biomed_llama 7b

app.py

@@ -190,28 +190,24 @@ 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
+    try:
+        model = AutoModel.from_pretrained(
+            "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
+            device_map="auto",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
+            trust_remote_code=True
+        )
+        
+        return model, tokenizer
+    except Exception as e:
+        print(f"Failed to initialize LLM: {str(e)}")
+        return None, None
 
 model = initialize_model()
 llm_model, llm_tokenizer = initialize_llm()
@@ -254,29 +250,37 @@ def process_image(image_path, text_prompts, modality):
             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
-        ])
+        # Process with LLM only if available
+        if llm_model is not None and llm_tokenizer is not None:
+            try:
+                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
+                ])
+            except Exception as e:
+                print(f"LLM analysis failed: {str(e)}")
+                combined_analysis = "\n".join(analysis_results)
+        else:
+            # If LLM is not available, only show BiomedParse results
+            combined_analysis = "\n".join(analysis_results)
         
         return results, combined_analysis
 

requirements.txt

@@ -52,3 +52,13 @@ pydot==3.0.1
 tabulate==0.9.0
 termcolor==2.4.0
 tokenizers==0.14.1
+# 
+# torch>=2.0.0
+# transformers>=4.34.0
+# gradio>=4.40.0
+# Pillow>=9.0.0
+# numpy>=1.21.0
+tqdm>=4.65.0
+# huggingface-hub>=0.19.0
+safetensors>=0.4.0
+# accelerate>=0.25.0