Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 from PIL import Image
 import os
 import gc
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import AutoProcessor, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel
 
 # Page config
@@ -37,7 +37,7 @@ device = init_device()
 
 @st.cache_resource
 def load_model():
-    """Load model with fallback options for quantization"""
+    """Load model with proper quantization handling"""
     try:
         # Using your original base model
         base_model_id = "unsloth/llama-3.2-11b-vision-instruct-unsloth-bnb-4bit"
@@ -45,22 +45,19 @@ def load_model():
         # Load processor
         processor = AutoProcessor.from_pretrained(base_model_id)
         
-        # Try to load with 4-bit quantization first
-        try:
-            import bitsandbytes
-            model = AutoModelForCausalLM.from_pretrained(
-                base_model_id,
-                device_map="auto",
+        # Load the model with proper quantization settings
+        model = AutoModelForCausalLM.from_pretrained(
+            base_model_id,
+            device_map="auto",
+            load_in_4bit=True,
+            quantization_config=BitsAndBytesConfig(
                 load_in_4bit=True,
-                torch_dtype=torch.float16
-            )
-        except ImportError:
-            st.warning("bitsandbytes not available. Falling back to float16 precision.")
-            model = AutoModelForCausalLM.from_pretrained(
-                base_model_id,
-                device_map="auto",
-                torch_dtype=torch.float16
-            )
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4"
+            ),
+            torch_dtype=torch.float16
+        )
         
         # Load adapter
         adapter_id = "saakshigupta/deepfake-explainer-1"