Loading model async

app.py

@@ -1,6 +1,8 @@
+import asyncio
 import io
 import logging
 import os
+import threading
 import uuid
 
 import streamlit as st
@@ -15,34 +17,68 @@ from transformers.image_utils import load_image
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Capture logs
-log_stream = io.StringIO()
-logging.basicConfig(stream=log_stream, level=logging.INFO)
-
+#log_stream = io.StringIO()
+#logging.basicConfig(stream=log_stream, level=logging.INFO)
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 if "session_id" not in st.session_state:
     st.session_state["session_id"] = str(uuid.uuid4())  # Generate unique session ID
 
 
-@st.cache_resource  # Streamlit Caching decorator
+# Async function to load the model
+async def load_model_embedding_async():
+    st.session_state["loading_model_embedding"] = True  # Show loading status
+    await asyncio.sleep(0.1)  # Allow UI updates
+    model_embedding = RAGMultiModalModel.from_pretrained("vidore/colpali-v1.2")
+    st.session_state["model_embedding"] = model_embedding
+    st.session_state["loading_model_embedding"] = False  # Model is ready
+
+
+# Function to run async function in a separate thread
 def load_model_embedding():
-    #docs_retrieval_model = RAGMultiModalModel.from_pretrained("vidore/colsmolvlm-alpha")
-    #docs_retrieval_model = RAGMultiModalModel.from_pretrained("vidore/colqwen2-v1.0")
-    docs_retrieval_model = RAGMultiModalModel.from_pretrained("vidore/colpali-v1.2")
-    return docs_retrieval_model
-model_embedding = load_model_embedding()
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(load_model_embedding_async())
 
-@st.cache_resource  # Streamlit Caching decorator
-def load_model_vlm():
+# Start model loading in a background thread
+if "model_embedding" not in st.session_state:
+    with st.status("Loading embedding model... ⏳"):
+        threading.Thread(target=load_model_embedding, daemon=True).start()
+
+
+
+# Async function to load the model
+async def load_model_vlm_async():
+    st.session_state["loading_model_vlm"] = True  # Show loading status
+    await asyncio.sleep(0.1)  # Allow UI updates
+    
     checkpoint = "HuggingFaceTB/SmolVLM-Instruct"
-    processor = AutoProcessor.from_pretrained(checkpoint)
+    processor_vlm = AutoProcessor.from_pretrained(checkpoint)
     quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-    model = AutoModelForVision2Seq.from_pretrained(
+    model_vlm = AutoModelForVision2Seq.from_pretrained(
         checkpoint,
         #torch_dtype=torch.bfloat16,
         quantization_config=quantization_config,
     )
-    return model, processor
-model_vlm, processor_vlm = load_model_vlm()
+
+    st.session_state["model_vlm"] = model_vlm
+    st.session_state["processor_vlm"] = processor_vlm
+    st.session_state["loading_model_vlm"] = False  # Model is ready
+
+
+# Function to run async function in a separate thread
+def load_model_vlm():
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(load_model_vlm_async())
+
+
+# Start model loading in a background thread
+if "model_vlm" not in st.session_state:
+    with st.status("Loading VLM model... ⏳"):
+        threading.Thread(target=load_model_vlm, daemon=True).start()
+
 
 
 
@@ -64,7 +100,7 @@ with st.sidebar:
     "[Source Code](https://huggingface.co/spaces/deepakkarkala/multimodal-rag/tree/main)"
 
 st.title("📝 Image Q&A with VLM")
-st.text_area("Logs:", log_stream.getvalue(), height=200)
+#st.text_area("Logs:", log_stream.getvalue(), height=200)
 
 uploaded_pdf = st.file_uploader("Upload PDF file", type=("pdf"))
 query = st.text_input(
@@ -73,16 +109,34 @@ query = st.text_input(
     disabled=not uploaded_pdf,
 )
 
+if st.session_state.get("loading_model_embedding", True):
+    st.warning("Loading Embedding model....")
+else:
+    st.success("Embedding Model loaded successfully! 🎉")
+
+if st.session_state.get("loading_model_vlm", True):
+    st.warning("Loading VLM model....")
+else:
+    st.success("VLM Model loaded successfully! 🎉")
+
+
+
+
+
+
+
 images = []
 images_folder = "data/" + st.session_state["session_id"] + "/"
 index_name = "index_" + st.session_state["session_id"]
 
 
-if uploaded_pdf and "is_index_complete" not in st.session_state:
+
+if uploaded_pdf and "model_embedding" in st.session_state and "is_index_complete" not in st.session_state:
     images = convert_from_bytes(uploaded_pdf.getvalue())
-    save_images_to_local(images, output_folder=images_folder)
+    save_images_to_local(images, output_folder=images_folder)    
+
     # index documents using the document retrieval model
-    model_embedding.index(
+    st.session_state["model_embedding"].index(
         input_path=images_folder, index_name=index_name, store_collection_with_index=False, overwrite=True
     )
     logging.info(f"{len(images)} number of images extracted from PDF and indexed")
@@ -90,13 +144,17 @@ if uploaded_pdf and "is_index_complete" not in st.session_state:
 
 
 
-if uploaded_pdf and query:
-    docs_retrieved = model_embedding.search(query, k=1)
+if uploaded_pdf and query and "model_embedding" in st.session_state and "model_vlm" in st.session_state:
+    docs_retrieved = st.session_state["model_embedding"].search(query, k=1)
     logging.info(f"{len(docs_retrieved)} number of images retrieved as relevant to query")
     image_id = docs_retrieved[0]["doc_id"]
     logging.info(f"Image id:{image_id} retrieved" )
     image_similar_to_query = images[image_id]
 
+
+    model_vlm, processor_vlm = st.session_state["model_vlm"], st.session_state["processor_vlm"]
+
+
     # Create input messages
     system_prompt = "You are an AI assistant. Your task is reply to user questions based on the provided image context."
     chat_template = [