Update app.py

app.py

@@ -403,17 +403,23 @@ def process_image_with_gradcam(image, model, device, pred_class):
 
 # ----- BLIP Image Captioning -----
 
-# Function to load BLIP captioning model
+# Function to load BLIP captioning models
 @st.cache_resource
-def load_blip_model():
-    with st.spinner("Loading BLIP captioning model..."):
+def load_blip_models():
+    with st.spinner("Loading BLIP captioning models..."):
         try:
-            processor = BlipProcessor.from_pretrained("saakshigupta/deepfake-blip-large")
-            model = BlipForConditionalGeneration.from_pretrained("saakshigupta/deepfake-blip-large")
-            return processor, model
+            # Load original BLIP model for general image captioning
+            original_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+            original_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
+            
+            # Load fine-tuned BLIP model for GradCAM analysis
+            finetuned_processor = BlipProcessor.from_pretrained("saakshigupta/deepfake-blip-large")
+            finetuned_model = BlipForConditionalGeneration.from_pretrained("saakshigupta/deepfake-blip-large")
+            
+            return original_processor, original_model, finetuned_processor, finetuned_model
         except Exception as e:
-            st.error(f"Error loading BLIP model: {str(e)}")
-            return None, None
+            st.error(f"Error loading BLIP models: {str(e)}")
+            return None, None, None, None
 
 # Function to generate image caption using BLIP's VQA approach for GradCAM
 def generate_gradcam_caption(image, processor, model, max_length=60):
@@ -452,7 +458,7 @@ def generate_gradcam_caption(image, processor, model, max_length=60):
 
 # Function to generate caption for original image
 def generate_image_caption(image, processor, model, max_length=75, num_beams=5):
-    """Generate a caption for the original image using BLIP model"""
+    """Generate a caption for the original image using the original BLIP model"""
     try:
         # Check for available GPU
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -594,8 +600,10 @@ def main():
     
     if 'blip_model_loaded' not in st.session_state:
         st.session_state.blip_model_loaded = False
-        st.session_state.blip_processor = None
-        st.session_state.blip_model = None
+        st.session_state.original_processor = None
+        st.session_state.original_model = None
+        st.session_state.finetuned_processor = None
+        st.session_state.finetuned_model = None
     
     # Initialize chat history
     if 'chat_history' not in st.session_state:
@@ -625,17 +633,19 @@ def main():
         with blip_col:
             if not st.session_state.blip_model_loaded:
                 if st.button("📥 Load BLIP for Captioning", type="primary"):
-                    # Load BLIP model
-                    processor, model = load_blip_model()
-                    if model is not None and processor is not None:
-                        st.session_state.blip_processor = processor
-                        st.session_state.blip_model = model
+                    # Load BLIP models
+                    original_processor, original_model, finetuned_processor, finetuned_model = load_blip_models()
+                    if all([original_processor, original_model, finetuned_processor, finetuned_model]):
+                        st.session_state.original_processor = original_processor
+                        st.session_state.original_model = original_model
+                        st.session_state.finetuned_processor = finetuned_processor
+                        st.session_state.finetuned_model = finetuned_model
                         st.session_state.blip_model_loaded = True
-                        st.success("✅ BLIP captioning model loaded successfully!")
+                        st.success("✅ BLIP captioning models loaded successfully!")
                     else:
-                        st.error("❌ Failed to load BLIP model.")
+                        st.error("❌ Failed to load BLIP models.")
             else:
-                st.success("✅ BLIP captioning model loaded and ready!")
+                st.success("✅ BLIP captioning models loaded and ready!")
                 
         with llm_col:
             if not st.session_state.llm_model_loaded:
@@ -672,8 +682,8 @@ def main():
                     with st.spinner("Generating image description..."):
                         caption = generate_image_caption(
                             image, 
-                            st.session_state.blip_processor, 
-                            st.session_state.blip_model
+                            st.session_state.original_processor, 
+                            st.session_state.original_model
                         )
                         st.session_state.image_caption = caption
                         
@@ -729,8 +739,8 @@ def main():
                             with st.spinner("Analyzing GradCAM visualization..."):
                                 gradcam_caption = generate_gradcam_caption(
                                     overlay,
-                                    st.session_state.blip_processor,
-                                    st.session_state.blip_model
+                                    st.session_state.finetuned_processor,
+                                    st.session_state.finetuned_model
                                 )
                                 st.session_state.gradcam_caption = gradcam_caption