Update app.py

app.py

@@ -16,20 +16,24 @@ st.set_page_config(page_title="Niche AI", layout="centered")
 st.title("🧠 Niche AI")
 st.caption("llama.cpp · CPU · Embedded · Streaming")
 
-@st.cache_resource
-def load_llm():
-    return Llama(
-        model_path=MODEL_PATH,
-        n_ctx=N_CTX,
-        n_threads=N_THREADS,
-        n_batch=N_BATCH,
-        f16_kv=True,
-        use_mmap=True,
-        use_mlock=False,
-        verbose=False,
-    )
-
-llm = load_llm()
+# ---------- LAZY LOADING ----------
+if "llm" not in st.session_state:
+    st.session_state.llm = None
+
+def get_llm():
+    if st.session_state.llm is None:
+        with st.spinner("Loading model..."):
+            st.session_state.llm = Llama(
+                model_path=MODEL_PATH,
+                n_ctx=N_CTX,
+                n_threads=N_THREADS,
+                n_batch=N_BATCH,
+                f16_kv=True,
+                use_mmap=True,
+                use_mlock=False,
+                verbose=False,
+            )
+    return st.session_state.llm
 
 # ---------- SESSION STATE ----------
 if "history" not in st.session_state:
@@ -94,6 +98,7 @@ def build_prompt(user_text: str) -> str:
 if st.button("Send") and prompt.strip():
     st.session_state.history.append(("You", prompt))
 
+    llm = get_llm()  # Lazy load here
     full_prompt = build_prompt(prompt)
 
     placeholder = st.empty()