update layout

app.py

@@ -85,10 +85,6 @@ except Exception as e:
 # pipe.enable_xformers_memory_efficient_attention()
 pipe.enable_attention_slicing()
 
-mesh_seqs = []
-frame_seqs = []
-cur_mesh_idx = None
-
 def read_video_frames(video_path, process_length, max_res):
     print("==> processing video: ", video_path)
     vid = VideoReader(video_path, ctx=cpu(0))
@@ -131,12 +127,10 @@ def infer_geometry(
     decode_chunk_size: int,
     overlap: int,
     downsample_ratio: float = 1.0, # downsample pcd for visualization
-    num_sample_frames: int =8, # downsample frames for visualization
     remove_edge: bool = True, # remove edge for visualization
     save_folder: str = os.path.join('workspace', 'GeometryCrafterApp'),
 ):
     try:
-        global cur_mesh_idx, mesh_seqs, frame_seqs
         run_id = str(uuid.uuid4())
         set_seed(42)
         pipe.enable_xformers_memory_efficient_attention()
@@ -206,11 +200,10 @@ def infer_geometry(
                 edge_mask = compute_edge_mask(point_maps[i, :, :, 2], 3)
                 valid_masks[i] = valid_masks[i] & (~edge_mask)
 
-        indices = np.linspace(0, len(point_maps)-1, num_sample_frames)
+        indices = np.linspace(0, len(point_maps)-1, 6)
         indices = np.round(indices).astype(np.int32)
-        
-        mesh_seqs.clear()
-        cur_mesh_idx = None
+
+        mesh_seqs, frame_seqs = [], []
 
         for index in indices:
 
@@ -224,20 +217,17 @@ def infer_geometry(
             mesh_seqs.append(output_glb_path)
             frame_seqs.append(index)
         
-        cur_mesh_idx = 0
         
         gc.collect()
         torch.cuda.empty_cache()
         
         return [
-            gr.Model3D(value=mesh_seqs[cur_mesh_idx], label=f"Frame: {frame_seqs[cur_mesh_idx]}"), 
+            gr.Model3D(value=mesh_seqs[idx], label=f"Frame: {frame_seqs[idx]}") for idx in range(len(frame_seqs))
+        ] + [ 
             gr.Video(value=output_disp_path, label="Disparity", interactive=False), 
             gr.DownloadButton("Download Npz File", value=output_npz_path, visible=True)
         ]
     except Exception as e:
-        mesh_seqs.clear()
-        frame_seqs.clear()
-        cur_mesh_idx = None
         gc.collect()
         torch.cuda.empty_cache()
         raise gr.Error(str(e))
@@ -251,21 +241,6 @@ def infer_geometry(
         #     gr.DownloadButton("Download Npz File", visible=False)
         # ]
 
-def goto_prev_frame():
-    global cur_mesh_idx, mesh_seqs, frame_seqs
-    if cur_mesh_idx is not None and len(mesh_seqs) > 0:
-        if cur_mesh_idx > 0:
-            cur_mesh_idx -= 1
-        return gr.Model3D(value=mesh_seqs[cur_mesh_idx], label=f"Frame: {frame_seqs[cur_mesh_idx]}")
-        
-
-def goto_next_frame():
-    global cur_mesh_idx, mesh_seqs, frame_seqs
-    if cur_mesh_idx is not None and len(mesh_seqs) > 0:
-        if cur_mesh_idx < len(mesh_seqs)-1:
-            cur_mesh_idx += 1
-        return gr.Model3D(value=mesh_seqs[cur_mesh_idx], label=f"Frame: {frame_seqs[cur_mesh_idx]}")
-
 def download_file():
     return gr.DownloadButton(visible=False)
 
@@ -353,23 +328,58 @@ def build_demo():
                     generate_btn = gr.Button("Generate")
 
             with gr.Column(scale=1):
-                output_point_maps = gr.Model3D(
-                    label="Point Map",
+                output_disp_video = gr.Video(
+                    label="Disparity",
+                    interactive=False
+                )
+                download_btn = gr.DownloadButton("Download Npz File", visible=False)
+
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=1):
+                output_point_map0 = gr.Model3D(
+                    label="Point Map 0",
                     clear_color=[1.0, 1.0, 1.0, 1.0],
                     # display_mode="solid"
                     interactive=False
                 )
-                with gr.Row():
-                    prev_btn = gr.Button("Prev")
-                    next_btn = gr.Button("Next")
-
             with gr.Column(scale=1):
-                output_disp_video = gr.Video(
-                    label="Disparity",
+                output_point_map1 = gr.Model3D(
+                    label="Point Map 1",
+                    clear_color=[1.0, 1.0, 1.0, 1.0],
+                    # display_mode="solid"
                     interactive=False
                 )
-                download_btn = gr.DownloadButton("Download Npz File", visible=False)
-
+            with gr.Column(scale=1):
+                output_point_map2 = gr.Model3D(
+                    label="Point Map 2",
+                    clear_color=[1.0, 1.0, 1.0, 1.0],
+                    # display_mode="solid"
+                    interactive=False
+                )
+        
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=1):
+                output_point_map3 = gr.Model3D(
+                    label="Point Map 3",
+                    clear_color=[1.0, 1.0, 1.0, 1.0],
+                    # display_mode="solid"
+                    interactive=False
+                )
+            with gr.Column(scale=1):
+                output_point_map4 = gr.Model3D(
+                    label="Point Map 4",
+                    clear_color=[1.0, 1.0, 1.0, 1.0],
+                    # display_mode="solid"
+                    interactive=False
+                )
+            with gr.Column(scale=1):
+                output_point_map5 = gr.Model3D(
+                    label="Point Map 5",
+                    clear_color=[1.0, 1.0, 1.0, 1.0],
+                    # display_mode="solid"
+                    interactive=False
+                )
+        
         gr.Examples(
             examples=examples,
             fn=infer_geometry,
@@ -383,7 +393,11 @@ def build_demo():
                 decode_chunk_size,
                 overlap,
             ],
-            outputs=[output_point_maps, output_disp_video, download_btn],            
+            outputs=[
+                output_point_map0, output_point_map1, output_point_map2,
+                output_point_map3, output_point_map4, output_point_map5,
+                output_disp_video, download_btn
+            ],            
             # cache_examples="lazy",
         )
         gr.Markdown(
@@ -412,17 +426,13 @@ def build_demo():
                 decode_chunk_size,
                 overlap,
             ],
-            outputs=[output_point_maps, output_disp_video, download_btn],
+            outputs=[
+                output_point_map0, output_point_map1, output_point_map2,
+                output_point_map3, output_point_map4, output_point_map5,
+                output_disp_video, download_btn
+            ],
         )
 
-        prev_btn.click(
-            fn=goto_prev_frame,
-            outputs=output_point_maps,
-        )
-        next_btn.click(
-            fn=goto_next_frame,
-            outputs=output_point_maps,
-        )
         download_btn.click(
             fn=download_file,
             outputs=download_btn