Update app.py

app.py

@@ -7,7 +7,8 @@ from PIL import Image
 from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 import gradio as gr
 
-# Load the model and tokenizer
+
+# Configuration for 4-bit quantization
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
@@ -15,20 +16,22 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16,
 )
 
+# Load the model without flash-attn
 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",
+    attn_implementation=None,  # Disable flash-attn
 )
 
 tokenizer = AutoTokenizer.from_pretrained(
-    "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1", 
+    "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
     trust_remote_code=True
 )
 
+
 # Define the function to handle the input
 def process_input(image, question):
     image = Image.open(image).convert("RGB")