Improve waiting behaviour in app during optimization

Whitebox_style_transfer.py

@@ -140,20 +140,27 @@ def img_choice_panel(imgtype, urls, default_choice, expanded):
         st.sidebar.write(f'Selected {imgtype} image:')
         st.sidebar.markdown(f'<img src="{state[f"{imgtype}_render_src"]}" width=240px></img>', unsafe_allow_html=True)
 
-
 def optimize(effect, preset, result_image_placeholder):
     content = st.session_state["Content_im"]
     style = st.session_state["Style_im"]
+    st.session_state["optimize_next"] = False
+    with st.spinner(text="Optimizing parameters.."):
+        if HUGGING_FACE:
+            optimize_on_server(content, style, result_image_placeholder)
+        else:
+            optimize_params(effect, preset, content, style, result_image_placeholder)
+
+def optimize_next(result_image_placeholder):
     result_image_placeholder.text("<- Custom content/style needs to be style transferred")
-    st.sidebar.warning("Note: Optimizing takes up to 5 minutes.")
+    queue_length = 0 if not HUGGING_FACE else get_queue_length()
+    if queue_length > 0:
+        st.sidebar.warning(f"WARNING: Already {queue_length} tasks in the queue. It will take approx {(queue_length+1) * 5} min for your image to be completed.")
+    else:
+        st.sidebar.warning("Note: Optimizing takes up to 5 minutes.")
     optimize_button = st.sidebar.button("Optimize Style Transfer")
     if optimize_button:
-        with st.spinner(text="Optimizing parameters.."):
-            if HUGGING_FACE:
-                optimize_on_server(content, style, result_image_placeholder)
-            else:
-                optimize_params(effect, preset, content, style, result_image_placeholder)
-        return st.session_state["effect_input"], st.session_state["result_vp"]
+        st.session_state["optimize_next"] = True
+        st.experimental_rerun()
     else:
         if not "result_vp" in st.session_state:
             st.stop()
@@ -206,15 +213,24 @@ coll2.header("Global Edits")
 result_image_placeholder = coll1.empty()
 result_image_placeholder.markdown("## loading..")
 
-from tasks import optimize_on_server, optimize_params, monitor_task
+from tasks import optimize_on_server, optimize_params, monitor_task, get_queue_length
 
 if "current_server_task_id" not in st.session_state:
     st.session_state['current_server_task_id'] = None
 
+if "optimize_next" not in st.session_state:
+    st.session_state['optimize_next'] = False
+
+effect, preset = create_effect()
+
 if HUGGING_FACE and st.session_state['current_server_task_id'] is not None: 
     with st.spinner(text="Optimizing parameters.."):
         monitor_task(result_image_placeholder)
 
+if st.session_state["optimize_next"]:
+    print("optimize now")
+    optimize(effect, preset, result_image_placeholder)
+
 img_choice_panel("Content", content_urls, "portrait", expanded=True)
 img_choice_panel("Style", style_urls, "starry_night", expanded=True)
 
@@ -222,11 +238,10 @@ state = session_state.get()
 content_id = state["Content_id"]
 style_id = state["Style_id"]
 
-effect, preset = create_effect()
 
 print("content id, style id", content_id, style_id  )
 if st.session_state["action"] == "uploaded":
-    content_img, _vp = optimize(effect, preset, result_image_placeholder)
+    content_img, _vp = optimize_next(result_image_placeholder)
 elif st.session_state["action"] in ("switch_page_from_local_edits", "switch_page_from_presets", "slider_change") or \
       content_id == "uploaded" or style_id == "uploaded":
     print("restore param")

pages/1_🎨_Apply_preset.py

@@ -42,7 +42,7 @@ presets = {
 
 st.session_state["action"] = "switch_page_from_presets" # on switchback, remember effect input
 
-active_preset = st.sidebar.selectbox("apply preset: ", ["original", "bump mapped", "contoured"])
+active_preset = st.sidebar.selectbox("apply preset: ", ["bump mapped", "contoured", "original"])
 blend_strength = st.sidebar.slider("Parameter blending strength (non-hue) : ", 0.0, 1.0, 1.0, 0.05)
 hue_blend_strength = st.sidebar.slider("Hue-shift blending strength : ", 0.0, 1.0, 1.0, 0.05)
 
@@ -118,4 +118,6 @@ coll2.image(img_res)
 
 apply_btn = st.sidebar.button("Apply")
 if apply_btn:
-    st.session_state["result_vp"] = vp
+    st.session_state["result_vp"] = vp
+
+st.info("Note: Press apply to make changes permanent")

tasks.py

@@ -53,55 +53,63 @@ def monitor_task(progress_placeholder):
 
     started_time = time.time()
     retries = 3
-    while True:
-        status = requests.get(WORKER_URL+"/get_status", params={"task_id": task_id})
-        if status.status_code != 200:
-            print("get_status got status_code", status.status_code)
-            st.warning(status.content)
-            retries -= 1
-            if retries == 0:
+    with progress_placeholder.container():
+        st.warning("Do not interact with the app until results are shown - otherwise results might be lost.")
+        progress_bar = st.empty()
+        while True:
+            status = requests.get(WORKER_URL+"/get_status", params={"task_id": task_id})
+            if status.status_code != 200:
+                print("get_status got status_code", status.status_code)
+                st.warning(status.content)
+                retries -= 1
+                if retries == 0:
+                    return
+                else:
+                    time.sleep(2)
+                    continue
+            status = status.json()
+            print(status)
+            if status["status"] != "running" and status["status"] != "queued" :
+                if status["msg"] != "":
+                    print("got error for task", task_id, ":", status["msg"])
+                    progress_placeholder.error(status["msg"])
+                    st.session_state['current_server_task_id'] = None
+                    st.stop()
+                if status["status"] == "finished":
+                    retrieve_for_results_from_server()
                 return
+            elif status["status"] == "queued":
+                started_time = time.time()
+                queue_length = requests.get(WORKER_URL+"/queue_length").json()
+                progress_bar.write(f"There are {queue_length['length']} tasks in the queue")
+            elif status["progress"] == 0.0:
+                progressed = min(0.5 * (time.time() - started_time) / 80.0, 0.5) #estimate 80s for strotts
+                progress_bar.progress(progressed)
             else:
-                time.sleep(2)
-                continue
-        status = status.json()
-        print(status)
-        if status["status"] != "running" and status["status"] != "queued" :
-            if status["msg"] != "":
-                print("got error for task", task_id, ":", status["msg"])
-                progress_placeholder.error(status["msg"])
-                st.session_state['current_server_task_id'] = None
-                st.stop()
-            if status["status"] == "finished":
-                retrieve_for_results_from_server()
-            return
-        elif status["status"] == "queued":
-            started_time = time.time()
-            queue_length = requests.get(WORKER_URL+"/queue_length").json()
-            progress_placeholder.write(f"There are {queue_length['length']} tasks in the queue")
-        elif status["progress"] == 0.0:
-            progressed = min(0.5 * (time.time() - started_time) / 80.0, 0.5) #estimate 80s for strotts
-            progress_placeholder.progress(progressed)
-        else:
-            progress_placeholder.progress(min(0.5 + status["progress"] / 2.0, 1.0))
-
-        time.sleep(2)
+                progress_bar.progress(min(0.5 + status["progress"] / 2.0, 1.0))
+
+            time.sleep(2)
+
+def get_queue_length():
+    queue_length = requests.get(WORKER_URL+"/queue_length").json()
+    return queue_length['length']
 
 
 def optimize_on_server(content, style, result_image_placeholder):
-    url = WORKER_URL + "/upload"
     content_path=f"/tmp/content-wise-uploaded{str(datetime.datetime.timestamp(datetime.datetime.now()))}.jpg"
     style_path=f"/tmp/content-wise-uploaded{str(datetime.datetime.timestamp(datetime.datetime.now()))}.jpg"
     asp_c, asp_s =  content.height / content.width, style.height / style.width
     if any(a < 0.5 or a > 2.0 for a in (asp_c, asp_s)):
         result_image_placeholder.error('aspect ratio must be <= 2')
         st.stop()
+
     content = pil_resize_long_edge_to(content, 1024)
     content.save(content_path)
     style = pil_resize_long_edge_to(style, 1024)
     style.save(style_path)
     files = {'style-image': open(style_path, "rb"), "content-image": open(content_path, "rb")}
     print("start-optimizing")
+    url = WORKER_URL + "/upload"
     task_id_res = requests.post(url, files=files)
     if task_id_res.status_code != 200:
         result_image_placeholder.error(task_id_res.content)