Update app.py

app.py

@@ -42,20 +42,25 @@ if upload_file is not None:
     image = Image.open(upload_file)
     st.image(image, caption="Uploaded Image", use_column_width=True)
 
+# Initialize session state for storing the generated code
+if "generated_code" not in st.session_state:
+    st.session_state.generated_code = None
+
 submit = st.button('Create a Webpage')
 if submit:
     image_data = input_image_setup(uploaded_file=upload_file)
     with st.spinner("Building the Webpage..."):
         response = generate_response(input_prompt, image_data)
         
-        # Extract code from the response
-        code = response.strip("```")
+        # Extract code from the response and store it in session state
+        st.session_state.generated_code = response.strip("```")
         
         # Display the code
         st.subheader("Generated Code:")
-        st.code(code, language='html')
-        
-    # Render the HTML Output with larger dimensions
+        st.code(st.session_state.generated_code, language='html')
+
+# Render the HTML Output with larger dimensions if code exists
+if st.session_state.generated_code:
     if st.button("Render This Code"):
         st.subheader("Rendered HTML Output:")
-        st.components.v1.html(code, height=1000, width=1500, scrolling=True)
+        st.components.v1.html(st.session_state.generated_code, height=1000, width=1500, scrolling=True)