Update app.py

app.py

@@ -38,9 +38,11 @@ st.warning("⚠️ Disclaimer: This app is currently running on CPU, which may r
 # Add download link
 st.markdown("[📥 Download the app code](https://huggingface.co/spaces/clayton07/colpali-qwen2-ocr/blob/main/app.py)")
 
-# Initialize session state for tracking if index is created
+# Initialize session state
 if 'index_created' not in st.session_state:
     st.session_state.index_created = False
+if 'processed_images' not in st.session_state:
+    st.session_state.processed_images = set()
 
 # File uploader
 image_source = st.radio("Choose image source:", ("Upload an image", "Use example image"))
@@ -61,23 +63,31 @@ if uploaded_file is not None:
     else:
         image_path = uploaded_file
 
-    if not st.session_state.index_created:
-        # Initialize the index for the first image
-        RAG.index(
-            input_path=image_path,
-            index_name="temp_index",
-            store_collection_with_index=False,
-            overwrite=True
-        )
-        st.session_state.index_created = True
-    else:
-        # Add to the existing index for subsequent images
-        RAG.add_to_index(
-            input_item=image_path,
-            store_collection_with_index=False
-        )
-
-    st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+    # Check if this image has been processed before
+    if image_path not in st.session_state.processed_images:
+        with st.spinner('Processing image...'):
+            if not st.session_state.index_created:
+                # Initialize the index for the first image
+                RAG.index(
+                    input_path=image_path,
+                    index_name="temp_index",
+                    store_collection_with_index=False,
+                    overwrite=True
+                )
+                st.session_state.index_created = True
+                st.success('Index created successfully!')
+            else:
+                # Add to the existing index for subsequent images
+                RAG.add_to_index(
+                    input_item=image_path,
+                    store_collection_with_index=False
+                )
+                st.success('Image added to index successfully!')
+            
+            # Mark this image as processed
+            st.session_state.processed_images.add(image_path)
+
+    st.image(image_path, caption="Uploaded Image", use_column_width=True)
 
     # Text query input
     text_query = st.text_input("Enter your query about the image:")
@@ -115,7 +125,7 @@ if uploaded_file is not None:
                 return_tensors="pt",
             )
             inputs = inputs.to(device)
-            generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens)  # Using the slider value here
+            generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens)
             generated_ids_trimmed = [
                 out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
             ]
@@ -128,7 +138,7 @@ if uploaded_file is not None:
         st.write(output_text[0])
 
     # Clean up temporary file
-    if image_source == "Upload an image":
+    if image_source == "Upload an image" and os.path.exists("temp_image.png"):
         os.remove("temp_image.png")
 else:
     st.write("Please upload an image to get started.")