Update app.py

app.py

@@ -1,9 +1,8 @@
 import streamlit as st
 import torch
 from PIL import Image
-from peft import PeftModel
-import gc
 import os
+import gc
 
 # Page config
 st.set_page_config(
@@ -39,7 +38,8 @@ device = init_device()
 def load_model():
     """Load model using Unsloth, similar to your notebook code"""
     try:
-        # Import Unsloth here to ensure it's loaded when needed
+        # Import libraries here to ensure they're loaded when needed
+        from peft import PeftModel
         from unsloth import FastVisionModel
         
         st.info("Loading base model and tokenizer using Unsloth...")
@@ -108,19 +108,20 @@ with st.sidebar:
     Model by [saakshigupta](https://huggingface.co/saakshigupta/deepfake-explainer-1)
     """)
 
-# Load model on startup
-with st.spinner("Loading model... this may take a minute."):
-    try:
-        model, tokenizer = load_model()
-        if model is not None and tokenizer is not None:
-            st.session_state['model'] = model
-            st.session_state['tokenizer'] = tokenizer
-            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)
+# Load model button
+if st.button("Load Model"):
+    with st.spinner("Loading model... this may take a minute."):
+        try:
+            model, tokenizer = load_model()
+            if model is not None and tokenizer is not None:
+                st.session_state['model'] = model
+                st.session_state['tokenizer'] = tokenizer
+                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"])
@@ -128,90 +129,92 @@ uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png
 # Check if model is 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:
+if uploaded_file is not None:
     # Display the image
     image = Image.open(uploaded_file).convert('RGB')
     st.image(image, caption="Uploaded Image", use_column_width=True)
     
-    # Analyze button
-    if st.button("Analyze Image"):
-        with st.spinner("Analyzing the image... This may take 15-30 seconds"):
-            try:
-                # Get components from session state
-                model = st.session_state['model']
-                tokenizer = st.session_state['tokenizer']
-                
-                # Format the message for Unsloth - same as your notebook
-                messages = [
-                    {"role": "user", "content": [
-                        {"type": "image"},
-                        {"type": "text", "text": custom_prompt}
-                    ]}
-                ]
-                
-                # Apply chat template
-                input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
-                
-                # Process with image
-                inputs = tokenizer(
-                    image,
-                    input_text,
-                    add_special_tokens=False,
-                    return_tensors="pt",
-                ).to(model.device)
-                
-                # Apply the cross-attention fix
-                fixed, inputs = fix_processor_outputs(inputs)
-                if fixed:
-                    st.info("Fixed cross-attention mask dimensions")
-                
-                # Generate analysis
-                with torch.no_grad():
-                    output_ids = model.generate(
-                        **inputs,
-                        max_new_tokens=max_length,
-                        temperature=temperature,
-                        top_p=0.9
-                    )
-                
-                # Decode the output
-                response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-                
-                # Extract the model's response
-                # Format might be different from processor.decode, check the output 
-                if "assistant" in response:
-                    result = response.split("assistant")[-1].strip()
-                else:
-                    result = response
-                
-                # Display result in a nice format
-                st.success("Analysis complete!")
-                
-                # Show technical and non-technical explanations separately if they exist
-                if "Technical Explanation:" in result and "Non-Technical Explanation:" in result:
-                    technical, non_technical = result.split("Non-Technical Explanation:")
-                    technical = technical.replace("Technical Explanation:", "").strip()
+    # Analyze button (only enabled if model is loaded)
+    if st.button("Analyze Image", disabled=not model_loaded):
+        if not model_loaded:
+            st.warning("Please load the model first by clicking the 'Load Model' button.")
+        else:
+            with st.spinner("Analyzing the image... This may take 15-30 seconds"):
+                try:
+                    # Get components from session state
+                    model = st.session_state['model']
+                    tokenizer = st.session_state['tokenizer']
+                    
+                    # Format the message for Unsloth - same as your notebook
+                    messages = [
+                        {"role": "user", "content": [
+                            {"type": "image"},
+                            {"type": "text", "text": custom_prompt}
+                        ]}
+                    ]
+                    
+                    # Apply chat template
+                    input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+                    
+                    # Process with image
+                    inputs = tokenizer(
+                        image,
+                        input_text,
+                        add_special_tokens=False,
+                        return_tensors="pt",
+                    ).to(model.device)
+                    
+                    # Apply the cross-attention fix
+                    fixed, inputs = fix_processor_outputs(inputs)
+                    if fixed:
+                        st.info("Fixed cross-attention mask dimensions")
+                    
+                    # Generate analysis
+                    with torch.no_grad():
+                        output_ids = model.generate(
+                            **inputs,
+                            max_new_tokens=max_length,
+                            temperature=temperature,
+                            top_p=0.9
+                        )
+                    
+                    # Decode the output
+                    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+                    
+                    # Extract the model's response
+                    if "assistant" in response:
+                        result = response.split("assistant")[-1].strip()
+                    else:
+                        result = response
+                    
+                    # Display result in a nice format
+                    st.success("Analysis complete!")
+                    
+                    # Show technical and non-technical explanations separately if they exist
+                    if "Technical Explanation:" in result and "Non-Technical Explanation:" in result:
+                        technical, non_technical = result.split("Non-Technical Explanation:")
+                        technical = technical.replace("Technical Explanation:", "").strip()
+                        
+                        col1, col2 = st.columns(2)
+                        with col1:
+                            st.subheader("Technical Analysis")
+                            st.write(technical)
+                        
+                        with col2:
+                            st.subheader("Simple Explanation")
+                            st.write(non_technical)
+                    else:
+                        st.subheader("Analysis Result")
+                        st.write(result)
                     
-                    col1, col2 = st.columns(2)
-                    with col1:
-                        st.subheader("Technical Analysis")
-                        st.write(technical)
+                    # Free memory after analysis
+                    free_memory()
                     
-                    with col2:
-                        st.subheader("Simple Explanation")
-                        st.write(non_technical)
-                else:
-                    st.subheader("Analysis Result")
-                    st.write(result)
-                
-                # Free memory after analysis
-                free_memory()
-                
-            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.")
+                except Exception as e:
+                    st.error(f"Error analyzing image: {str(e)}")
+                    st.exception(e)
+    elif not model_loaded:
+        st.warning("Please load the model first by clicking the 'Load Model' button at the top of the page.")
 else:
     st.info("Please upload an image to begin analysis")