Issue fix

app.py

@@ -20,48 +20,62 @@ nest_asyncio.apply()
 device='cuda' if torch.cuda.is_available() else 'cpu'
 st.set_page_config(page_title="DermBOT", page_icon="🧬", layout="centered")
 
+
+@st.cache_resource(show_spinner=False)
+def load_models():
+    """Cache all models to load only once"""
+    with st.spinner("Loading AI models (one-time operation)..."):
+        classifier = initialize_classifier()
+        return classifier
+
+@st.cache_resource(show_spinner=False)
+def initialize_llm(selected_model, api_key):
+    """Initialize the LLM based on selection"""
+    if "OpenAI" in selected_model:
+        return ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=api_key)
+    elif "LLaMA" in selected_model:
+        st.warning("LLaMA integration is not implemented yet.")
+        st.stop()
+    elif "Gemini" in selected_model:
+        st.warning("Gemini integration is not implemented yet.")
+        st.stop()
+    else:
+        st.error("Unsupported model selected.")
+        st.stop()
+
+@st.cache_resource(show_spinner=False)
+def load_rag_chain(llm):
+    """Initialize RAG chain only once"""
+    return invoke_rag_chain(llm)
+
 # === Model Selection ===
 available_models = ["OpenAI GPT-4o", "LLaMA 3", "Gemini Pro"]
-st.session_state["selected_model"] = st.sidebar.selectbox("Select LLM Model", available_models)
+if "selected_model" not in st.session_state:
+    st.session_state["selected_model"] = available_models[0]
+
+st.session_state["selected_model"] = st.sidebar.selectbox(
+    "Select LLM Model",
+    available_models,
+    index=available_models.index(st.session_state["selected_model"])
+)
 
 # Dynamically initialize LLM based on selection
 OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"]
-selected_model = st.session_state["selected_model"]
-if "OpenAI" in selected_model:
-    llm = ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=OPENAI_API_KEY)
-elif "LLaMA" in selected_model:
-    st.warning("LLaMA integration is not implemented yet.")
-    st.stop()
-elif "Gemini" in selected_model:
-    st.warning("Gemini integration is not implemented yet.")
-    st.stop()
-else:
-    st.error("Unsupported model selected.")
-    st.stop()
-
-
-with st.spinner("Loading AI models (one-time operation)..."):
-    classifier = initialize_classifier()
-st.success("Models loaded successfully!")
+llm = initialize_llm(st.session_state["selected_model"], OPENAI_API_KEY)
+rag_chain = load_rag_chain(llm)
+
+classifier = load_models()
 
 # === Session Init ===
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
+if "current_image" not in st.session_state:
+    st.session_state.current_image = None
 
 # === Image Processing Function ===
 def run_inference(image):
-
     result = classifier.predict(image, top_k=1)
-
-    # Display results
-    print("Top Predictions:")
-    for name, prob in result["top_predictions"]:
-        print(f"{name}: {prob:.2%}")
-
-    print("\nYou can also access all probabilities:")
-    print(result["all_probabilities"])
-
     predicted_label = result["top_predictions"][0][0]
     return predicted_label
 
@@ -82,43 +96,59 @@ def export_chat_to_pdf(messages):
 
 # === App UI ===
 
-print("=== Before App UI ===")
 st.title("🧬 DermBOT — Skin AI Assistant")
-st.caption(f"🧠 Using model: {selected_model}")
-uploaded_file = st.file_uploader("Upload a skin image", type=["jpg", "jpeg", "png"])
+st.caption(f"🧠 Using model: {st.session_state['selected_model']}")
+uploaded_file = st.file_uploader(
+    "Upload a skin image",
+    type=["jpg", "jpeg", "png"],
+    key="file_uploader"
+)
+
+if uploaded_file is not None and uploaded_file != st.session_state.current_image:
+    st.session_state.messages = []
+    st.session_state.current_image = uploaded_file
 
-if uploaded_file:
-    print("=== At start of uploading ===")
-    st.image(uploaded_file, caption="Uploaded image", use_column_width=True)
     image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Uploaded image", use_column_width=True)
+    with st.spinner("Analyzing the image..."):
+        predicted_label = run_inference(image)
 
-    predicted_label = run_inference(image)
-
-    # Show predictions clearly to the user
     st.markdown(f" Most Likely Diagnosis : {predicted_label}")
 
-    query = f"What are my treatment options for {predicted_label}?"
-    st.session_state.messages.append({"role": "user", "content": query})
+    initial_query = f"What are my treatment options for {predicted_label}?"
+    st.session_state.messages.append({"role": "user", "content": initial_query})
 
-    with st.spinner("Analyzing the image and retrieving response..."):
-        response = invoke_rag_chain(llm).invoke(query)
+    with st.spinner("Retrieving medical information..."):
+        response = rag_chain.invoke(initial_query)
         st.session_state.messages.append({"role": "assistant", "content": response['result']})
 
-    with st.chat_message("assistant"):
-        st.markdown(response['result'])
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
 # === Chat Interface ===
-if prompt := st.chat_input("Ask a follow-up..."):
+if prompt := st.chat_input("Ask a follow-up question..."):
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
 
-    response = llm.invoke([{"role": m["role"], "content": m["content"]} for m in st.session_state.messages])
-    st.session_state.messages.append({"role": "assistant", "content": response.content})
     with st.chat_message("assistant"):
-        st.markdown(response.content)
+        with st.spinner("Thinking..."):
+            if len(st.session_state.messages) > 1:
+                response = llm.invoke([{"role": m["role"], "content": m["content"]} for m in st.session_state.messages])
+            else:
+                response = rag_chain.invoke(prompt)
+                response = response['result']
+
+            st.markdown(response)
+            st.session_state.messages.append({"role": "assistant", "content": response})
+
 
-# === PDF Button ===
-if st.button("📄 Download Chat as PDF"):
+if st.session_state.messages and st.button("📄 Download Chat as PDF"):
     pdf_file = export_chat_to_pdf(st.session_state.messages)
-    st.download_button("Download PDF", data=pdf_file, file_name="chat_history.pdf", mime="application/pdf")
+    st.download_button(
+        "Download PDF",
+        data=pdf_file,
+        file_name="dermbot_chat_history.pdf",
+        mime="application/pdf"
+    )