Update app.py

app.py

@@ -5,6 +5,11 @@ from chatbot import generate_html_css_from_image
 def main():
     st.set_page_config(page_title="Gemini 2.5 Flash HTML/CSS Chatbot", layout="wide")
 
+    if "output_html" not in st.session_state:
+        st.session_state.output_html = None
+    if "last_image" not in st.session_state:
+        st.session_state.last_image = None
+
     st.markdown("""
         <div style='text-align: center; padding: 10px 0;'>
             <h1>Welcome, Gemini 2.5 Flash HTML/CSS Chatbot</h1>
@@ -47,29 +52,39 @@ def main():
 
     with right:
         if uploaded_image:
-            if "output_html" not in st.session_state or st.session_state.get("last_image") != uploaded_image.name:
+            current_name = getattr(uploaded_image, "name", "uploaded_image")
+
+            if st.session_state.last_image != current_name or st.session_state.output_html is None:
                 with st.spinner("Generating HTML + CSS using Gemini 2.5 Flash..."):
                     try:
                         output = generate_html_css_from_image(uploaded_image)
                         st.session_state.output_html = output
-                        st.session_state.last_image = uploaded_image.name
+                        st.session_state.last_image = current_name
                     except Exception as e:
                         st.error(f"Error: {e}")
                         return
 
-        output = st.session_state.output_html
-        if output:
-            st.subheader("💬 Generated HTML + CSS:")
-            placeholder = st.empty()
-            typed_text = ""
-            for char in output:
-                typed_text += char
-                placeholder.code(typed_text, language="html")
-                time.sleep(0.002)
+            output = st.session_state.get("output_html")
+            if output:
+                st.subheader("💬 Generated HTML + CSS:")
+                placeholder = st.empty()
+                typed_text = ""
+                for char in output:
+                    typed_text += char
+                    placeholder.code(typed_text, language="html")
+                    time.sleep(0.002)
 
-            st.download_button("💾 Download HTML file", data=output, file_name="generated_ui.html", mime="text/html")
+                st.download_button(
+                    "💾 Download HTML file",
+                    data=output,
+                    file_name="generated_ui.html",
+                    mime="text/html"
+                )
+            else:
+                st.error("No output generated.")
         else:
-            st.error("No output generated.")
+            st.session_state.output_html = None
+            st.session_state.last_image = None
 
 if __name__ == "__main__":
     main()