update UI

app.py

@@ -1,13 +1,13 @@
 # Author: Huzheng Yang
 # %%
-USE_SPACES = True
+HUGGINGFACE_SPACE = True
 BATCH_SIZE = 4
 
-if USE_SPACES:  # huggingface ZeroGPU, dynamic GPU allocation
+if HUGGINGFACE_SPACE:  # huggingface ZeroGPU, dynamic GPU allocation
     try:
         import spaces
     except ImportError:
-        USE_SPACES = False  # run on local machine
+        HUGGINGFACE_SPACE = False  # run on local machine
         BATCH_SIZE = 1
 
 import os
@@ -309,7 +309,7 @@ def _ncut_run(*args, **kwargs):
             torch.cuda.empty_cache()
         return [], "Error: " + str(e)
 
-if USE_SPACES:
+if HUGGINGFACE_SPACE:
     @spaces.GPU(duration=20)
     def quick_run(*args, **kwargs):
         return _ncut_run(*args, **kwargs)
@@ -326,7 +326,7 @@ if USE_SPACES:
     def super_duper_long_run(*args, **kwargs):
         return _ncut_run(*args, **kwargs)
 
-if not USE_SPACES:
+if not HUGGINGFACE_SPACE:
     def quick_run(*args, **kwargs):
         return _ncut_run(*args, **kwargs)
 
@@ -582,7 +582,7 @@ def make_parameters_section():
     from backbone import get_all_model_names
     model_names = get_all_model_names()
     model_dropdown = gr.Dropdown(model_names, label="Backbone", value="DiNO(dino_vitb8)", elem_id="model_name")
-    layer_slider = gr.Slider(1, 12, step=1, label="Backbone: Layer index", value=12, elem_id="layer")
+    layer_slider = gr.Slider(1, 12, step=1, label="Backbone: Layer index", value=10, elem_id="layer")
     node_type_dropdown = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?")
     num_eig_slider = gr.Slider(1, 1000, step=1, label="NCUT: Number of eigenvectors", value=100, elem_id="num_eig", info='increase for more clusters')
 
@@ -827,7 +827,7 @@ with gr.Blocks() as demo:
     with gr.Tab('Compare'): 
         def add_one_model(i_model=1):
             with gr.Column(scale=5, min_width=200) as col:
-                gr.Markdown(f'### Output Model {i_model}')
+                gr.Markdown(f'### Output Images')
                 output_gallery = gr.Gallery(value=[], label="NCUT Embedding", show_label=False, elem_id=f"ncut{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto")
                 submit_button = gr.Button("🔴RUN", elem_id=f"submit_button{i_model}")
                 [
@@ -864,29 +864,44 @@ with gr.Blocks() as demo:
                 load_images_button.click(lambda x: default_images, outputs=input_gallery)
 
                 
-            for i in range(1, 3):
-                add_one_model(i)
+            for i in range(2):
+                add_one_model()
                 
-        with gr.Row():
-            for i in range(1, 4):
-                with gr.Column(scale=5, min_width=200):
-                    slot = gr.Button("Add model", elem_id=f"add_model_button{i}")
-                    col = add_one_model(i+2)
-                    col.visible = False
-                    
-                    slot.click(
-                        fn=lambda x: gr.update(visible=True),
-                        outputs=col
-                    )
-                    slot.click(
-                        fn=lambda x: gr.update(visible=False),
-                        outputs=slot
-                    )
+        # Create rows and buttons in a loop
+        rows = []
+        buttons = []
 
+        for i in range(4):
+            row = gr.Row(visible=False)
+            rows.append(row)
+            
+            with row:
+                for j in range(3):
+                    with gr.Column(scale=5, min_width=200):
+                        add_one_model()
 
+            button = gr.Button("Add Compare", elem_id=f"add_button_{i}", visible=False if i > 0 else True)
+            buttons.append(button)
             
-download_all_models()
-download_all_datasets()
-demo.launch(share=True)
+            if i > 0:
+                # Reveal the current row and next button
+                buttons[i - 1].click(fn=lambda x: gr.update(visible=True), outputs=row)
+                buttons[i - 1].click(fn=lambda x: gr.update(visible=True), outputs=button)
+                
+                # Hide the current button
+                buttons[i - 1].click(fn=lambda x: gr.update(visible=False), outputs=buttons[i - 1])
+
+        # Last button only reveals the last row and hides itself
+        buttons[-1].click(fn=lambda x: gr.update(visible=True), outputs=rows[-1])
+        buttons[-1].click(fn=lambda x: gr.update(visible=False), outputs=buttons[-1])
+        
+
+
+if HUGGINGFACE_SPACE:
+    download_all_models()
+    download_all_datasets()
+    demo.launch()
+else:
+    demo.launch(share=True)
 
 # %%