make models singleton

app.py

@@ -37,13 +37,33 @@ text_embed_model = HuggingFaceEmbedding(
     embed_batch_size=1,
 )
 
+class IndexManager:
+    """Avoids deepcopying the index object in gr.State"""
+    def __init__(self):
+        self.current_index = None
+        # Initialize with default index
+        self.load_index(example_indexes["IONIQ 2024"])
+    
+    def load_index(self, index_path):
+        storage_context = StorageContext.from_defaults(persist_dir=index_path)
+        self.current_index = load_index_from_storage(
+            storage_context,
+            embed_model=text_embed_model,
+            image_embed_model=image_embed_model,
+        )
+        return f"Loaded index: {index_path}"
+    
+    def set_index(self, index):
+        self.current_index = index
+    
+    def get_index(self):
+        return self.current_index
+
+index_manager = IndexManager()
+
 def load_index(index_path: str) -> MultiModalVectorStoreIndex:
-    storage_context = StorageContext.from_defaults(persist_dir=index_path)
-    return load_index_from_storage(
-        storage_context,
-        embed_model=text_embed_model,
-        image_embed_model=image_embed_model,
-    )
+    index_manager.load_index(index_path)
+    return index_manager.get_index()
 
 def create_index(file, llama_parse_key, progress=gr.Progress()):
     if not file or not llama_parse_key:
@@ -88,14 +108,17 @@ def create_index(file, llama_parse_key, progress=gr.Progress()):
         )
         
         progress(1.0, desc="Complete!")
-        return index, "Index created successfully!"
+        index_manager.set_index(index)
+        return "Index created successfully!"
     
     except Exception as e:
-        return None, f"Error creating index: {str(e)}"
+        return f"Error creating index: {str(e)}"
 
-def run_search(index, query, text_top_k, image_top_k):
+def run_search(query, text_top_k, image_top_k):
+    index = index_manager.get_index()
     if not index:
         return "Please create or select an index first.", [], []
+    
     retriever = index.as_retriever(
         similarity_top_k=text_top_k,
         image_similarity_top_k=image_top_k,
@@ -194,41 +217,34 @@ Processing will take a few minutes when creating a new index, depending on the s
                 elem_id="text_results"
             )
     
-    # State
-    index_state = gr.State()
-
-    # Load default index on startup
-    default_index = load_index(example_indexes["IONIQ 2024"])
-    index_state.value = default_index
-    
     # Event handlers
     def load_existing_index(index_name):
         if index_name:
             try:
-                index = load_index(example_indexes[index_name])
-                return index, f"Loaded index: {index_name}"
+                index_manager.load_index(example_indexes[index_name])
+                return f"Loaded index: {index_name}"
             except Exception as e:
-                return None, f"Error loading index: {str(e)}"
-        return None, "No index selected"
+                return f"Error loading index: {str(e)}"
+        return "No index selected"
     
     existing_index_dropdown.change(
         fn=load_existing_index,
         inputs=[existing_index_dropdown],
-        outputs=[index_state, create_status],
+        outputs=[create_status],
         api_name=False
     )
     
     create_btn.click(
         fn=create_index,
         inputs=[file_upload, llama_parse_key],
-        outputs=[index_state, create_status],
+        outputs=[create_status],
         api_name=False,
-        show_progress=True  # Enable progress bar
+        show_progress=True
     )
     
     search_btn.click(
         fn=run_search,
-        inputs=[index_state, query_input, text_top_k, image_top_k],
+        inputs=[query_input, text_top_k, image_top_k],
         outputs=[status_output, text_output, image_output],
         api_name=False
     )