update UI

app.py

@@ -140,11 +140,13 @@ def image_to_tensor(img, width=256, height=256):
     return img
 
 
-@spaces.GPU
-def run_pose_exploration_a(image1, image2, seed_value):
+@spaces.GPU(duration=110)
+def run_pose_exploration(cam_vis, image1, image2, probe_bsz, adj_bsz, adj_iters, seed_value):
 
     seed_everything(seed_value)
 
+    cam_vis.set_images([np.asarray(image1, dtype=np.uint8), np.asarray(image2, dtype=np.uint8)])
+
     image1 = image_to_tensor(image1).to(_device_)
     image2 = image_to_tensor(image2).to(_device_)
 
@@ -156,22 +158,8 @@ def run_pose_exploration_a(image1, image2, seed_value):
         matcher_ckpt_path=_matcher_ckpt_path_
     )
 
-    return elevs, elev_ranges, gr.update(value='Preparation Done!')
-
-
-@spaces.GPU
-def run_pose_exploration_b(cam_vis, image1, image2, elevs, elev_ranges, probe_bsz, adj_bsz, adj_iters, seed_value):
-
-    seed_everything(seed_value)
-
     noise = np.random.randn(probe_bsz, 4, 32, 32)
 
-    cam_vis.set_images([np.asarray(image1, dtype=np.uint8), np.asarray(image2, dtype=np.uint8)])
-
-    image1 = image_to_tensor(image1).to(_device_)
-    image2 = image_to_tensor(image2).to(_device_)
-
-    images = [image1, image2]
     result_poses, aux_data = estimate_poses(
         _model_, images,
         seed_cand_num=8,
@@ -210,7 +198,7 @@ def run_pose_exploration_b(cam_vis, image1, image2, elevs, elev_ranges, probe_bs
     return anchor_polar, explored_sph, fig, gr.update(interactive=True)
 
 
-@spaces.GPU
+@spaces.GPU(duration=110)
 def run_pose_refinement(cam_vis, image1, image2, anchor_polar, explored_sph, refine_iters, seed_value):
 
     seed_everything(seed_value)
@@ -295,16 +283,13 @@ def run_demo():
                     with gr.Column(min_width=280):
                         processed_image2 = gr.Image(type='numpy', image_mode='RGB', label='Processed Image 2', width=280, interactive=False)
 
-                with gr.Row():
-                    progress_info = gr.Markdown(None, visible=False)
-
                 with gr.Row():
                     preprocess_chk = gr.Checkbox(True, label='Remove background and recenter object')
 
                 with gr.Accordion('Advanced options', open=False):
                     probe_bsz = gr.Slider(4, 32, value=16, step=4, label='Probe Batch Size')
                     adj_bsz = gr.Slider(1, 8, value=4, step=1, label='Adjust Batch Size')
-                    adj_iters = gr.Slider(1, 20, value=8, step=1, label='Adjust Iterations')
+                    adj_iters = gr.Slider(1, 20, value=10, step=1, label='Adjust Iterations')
                     seed_value = gr.Number(value=0, label="Seed Value", precision=0)
 
                 with gr.Row():
@@ -374,24 +359,14 @@ def run_demo():
         explored_sph = gr.State()
         anchor_polar = gr.State()
         refined_sph = gr.State()
-        elevs = gr.State()
-        elev_ranges = gr.State()
 
         run_btn.click(
             fn=run_preprocess,
             inputs=[input_image1, input_image2, preprocess_chk, seed_value],
             outputs=[processed_image1, processed_image2],
         ).success(
-            fn=lambda: gr.update(value='Preparing Estimation...', visible=True),
-            inputs=[],
-            outputs=[progress_info]
-        ).success(
-            fn=run_pose_exploration_a,
-            inputs=[processed_image1, processed_image2, seed_value],
-            outputs=[elevs, elev_ranges, progress_info]
-        ).success(
-            fn=partial(run_pose_exploration_b, cam_vis),
-            inputs=[processed_image1, processed_image2, elevs, elev_ranges, probe_bsz, adj_bsz, adj_iters, seed_value],
+            fn=partial(run_pose_exploration, cam_vis),
+            inputs=[processed_image1, processed_image2, probe_bsz, adj_bsz, adj_iters, seed_value],
             outputs=[anchor_polar, explored_sph, vis_output, refine_btn]
         )