Update app.py

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import torch
 import os
 from PIL import Image
-from transformers import AutoProcessor, MllamaForCausalLM, BitsAndBytesConfig
+from transformers import AutoProcessor, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel
 import gc
 
@@ -22,9 +22,8 @@ def free_memory():
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
 
-# Helper functions
+# Helper function to check CUDA
 def init_device():
     """Set the appropriate device and return it"""
     if torch.cuda.is_available():
@@ -39,13 +38,15 @@ device = init_device()
 
 @st.cache_resource
 def load_model():
-    """Load model and processor with caching to avoid reloading"""
+    """Load model and processor with proper dtype settings"""
     try:
         # Load base model
         base_model_id = "unsloth/llama-3.2-11b-vision-instruct-unsloth-bnb-4bit"
+        
+        # Load processor first
         processor = AutoProcessor.from_pretrained(base_model_id)
         
-        # Configure 4-bit quantization with correct dtype
+        # Configure quantization explicitly with float16
         quantization_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_compute_dtype=torch.float16,
@@ -53,11 +54,11 @@ def load_model():
             bnb_4bit_use_double_quant=True
         )
         
-        # Load model with explicit dtype settings using MllamaForCausalLM
-        model = MllamaForCausalLM.from_pretrained(
+        # Load model with explicit dtype settings
+        model = AutoModelForCausalLM.from_pretrained(
             base_model_id, 
             device_map="auto",
-            torch_dtype=torch.float16,
+            torch_dtype=torch.float16,  # Explicit float16
             quantization_config=quantization_config
         )
         
@@ -66,9 +67,10 @@ def load_model():
         model = PeftModel.from_pretrained(model, adapter_id)
         
         return model, processor
-            
+    
     except Exception as e:
         st.error(f"Error loading model: {str(e)}")
+        st.exception(e)
         return None, None
 
 # Function to fix cross-attention masks
@@ -112,34 +114,25 @@ with st.sidebar:
     Model by [saakshigupta](https://huggingface.co/saakshigupta/deepfake-explainer-1)
     """)
 
-# Load model on app startup with a progress bar
-if 'model_loaded' not in st.session_state:
-    progress_bar = st.progress(0)
-    st.info("Loading model... this may take a minute.")
-    
-    for i in range(10):
-        # Simulate progress while model loads
-        progress_bar.progress((i + 1) * 10)
-        if i == 2:
-            # Start loading the model at 30% progress
-            model, processor = load_model()
-            if model is not None:
-                st.session_state['model'] = model
-                st.session_state['processor'] = processor
-                st.session_state['model_loaded'] = True
-    
-    progress_bar.empty()
-    
-    if 'model_loaded' in st.session_state and st.session_state['model_loaded']:
-        st.success("Model loaded successfully!")
-    else:
-        st.error("Failed to load model. Try refreshing the page.")
+# Load model on startup
+with st.spinner("Loading model... this may take a minute."):
+    try:
+        model, processor = load_model()
+        if model is not None and processor is not None:
+            st.session_state['model'] = model
+            st.session_state['processor'] = processor
+            st.success("Model loaded successfully!")
+        else:
+            st.error("Failed to load model.")
+    except Exception as e:
+        st.error(f"Error during model loading: {str(e)}")
+        st.exception(e)
 
 # Main content area - file uploader
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
 # Check if model is loaded
-model_loaded = 'model_loaded' in st.session_state and st.session_state['model_loaded']
+model_loaded = 'model' in st.session_state and st.session_state['model'] is not None
 
 if uploaded_file is not None and model_loaded:
     # Display the image
@@ -208,6 +201,9 @@ if uploaded_file is not None and model_loaded:
                 
             except Exception as e:
                 st.error(f"Error analyzing image: {str(e)}")
+                st.exception(e)
+elif not model_loaded and uploaded_file is not None:
+    st.warning("Model not loaded correctly. Try refreshing the page.")
 else:
     st.info("Please upload an image to begin analysis")