Update app.py

app.py

@@ -38,11 +38,9 @@ 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
+# Initialize session state for tracking if index is created
 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"))
@@ -63,31 +61,23 @@ if uploaded_file is not None:
     else:
         image_path = uploaded_file
 
-    # 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)
+    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)
 
     # Text query input
     text_query = st.text_input("Enter your query about the image:")
@@ -125,7 +115,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)
+            generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens)  # Using the slider value here
             generated_ids_trimmed = [
                 out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
             ]
@@ -138,7 +128,7 @@ if uploaded_file is not None:
         st.write(output_text[0])
 
     # Clean up temporary file
-    if image_source == "Upload an image" and os.path.exists("temp_image.png"):
+    if image_source == "Upload an image":
         os.remove("temp_image.png")
 else:
     st.write("Please upload an image to get started.")