Update app.py

app.py

@@ -134,6 +134,12 @@ if prompt := st.chat_input("Ask your inspection question..."):
         st.error("🔑 Authentication required! Please enter your Hugging Face token in the sidebar.")
         st.stop()
     
+    # Load model if not loaded
+    if "model" not in st.session_state:
+        st.session_state.model, st.session_state.tokenizer = load_model(hf_token)
+    model = st.session_state.model
+    tokenizer = st.session_state.tokenizer
+    
     # Add user message with proper avatar handling
     try:
         with st.chat_message("user", avatar="👤"):
@@ -144,15 +150,22 @@ if prompt := st.chat_input("Ask your inspection question..."):
             st.markdown(prompt)
     
     st.session_state.messages.append({"role": "user", "content": prompt})
+
+    # Process file
+    file_context = process_file(uploaded_file) if uploaded_file else ""
     
-    # Generate and stream response
+    # Generate and stream response with avatar handling
     if model and tokenizer:
-        with st.chat_message("assistant", avatar="🤖"):
-            try:
+        try:
+            with st.chat_message("assistant", avatar="🤖"):
+                streamer = generate_response(prompt, file_context)
+                response = st.write_stream(streamer)
+                st.session_state.messages.append({"role": "assistant", "content": response})
+        except:
+            # Fallback if avatar fails
+            with st.chat_message("assistant"):
                 streamer = generate_response(prompt, file_context)
                 response = st.write_stream(streamer)
                 st.session_state.messages.append({"role": "assistant", "content": response})
-            except Exception as e:
-                st.error(f"⚡ Generation error: {str(e)}")
     else:
         st.error("🤖 Model not loaded - check your token and connection!")