Update app.py

app.py

@@ -3,7 +3,7 @@ import streamlit as st
 import torch
 from PIL import Image
 import io
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import AutoProcessor, BitsAndBytesConfig, MllamaForCausalLM
 from peft import PeftModel
 
 # Page config
@@ -21,12 +21,23 @@ st.markdown("Upload an image to analyze it for possible deepfake manipulation")
 def load_model():
     """Load model and processor (cached to avoid reloading)"""
     # Load base model
-    base_model_id = "unsloth/llama-3.2-11b-vision-instruct"
+    base_model_id = "unsloth/llama-3.2-11b-vision-instruct-unsloth-bnb-4bit"
     processor = AutoProcessor.from_pretrained(base_model_id)
-    model = AutoModelForCausalLM.from_pretrained(
+    
+    # Configure 4-bit quantization
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True
+    )
+    
+    # Use the specific model class MllamaForCausalLM instead of AutoModelForCausalLM
+    model = MllamaForCausalLM.from_pretrained(
         base_model_id, 
         device_map="auto", 
-        torch_dtype=torch.float16
+        torch_dtype=torch.float16,
+        quantization_config=quantization_config
     )
     
     # Load adapter