Update app.py

app.py

@@ -55,6 +55,7 @@ if not st.session_state["models_loaded"]:
 prompt = st.text_input("What do you want to see?")
 
 DEBUG = False
+# UI code taken from https://huggingface.co/spaces/flax-community/dalle-mini/blob/main/app/streamlit/app.py
 if prompt != "":
     container = st.empty()
     container.markdown(
@@ -74,7 +75,6 @@ if prompt != "":
         </div>
         </div>
         </div>
-        <small><i>Predictions may take up to 40s under high load. Please stand by.</i></small>
     """,
         unsafe_allow_html=True,
     )
@@ -90,3 +90,5 @@ if prompt != "":
     container.markdown(f"**{prompt}**")
 
     st.button("Again!", key="again_button")
+    
+     container.markdown("<small><i>Predictions may take up to 40s under high load. Please stand by.</i></small>")