Update id selector

app.py

@@ -34,7 +34,7 @@ image_size = 512
 
 @spaces.GPU  # Uncomment if using ZeroGPU
 def infer(
-    which_id,
+    selected_identity,
     background,
     pose,
     negative_prompt,
@@ -46,7 +46,7 @@ def infer(
     num_images=1
 ):
     
-    full_lora_weights_path = f"{folder_of_lora_weights}/{which_id}/{which_checkpoint}/{lora_name}"
+    full_lora_weights_path = f"{folder_of_lora_weights}/{selected_identity}/{which_checkpoint}/{lora_name}"
     pipe.load_lora_weights(full_lora_weights_path)
 
     if randomize_seed:
@@ -55,7 +55,7 @@ def infer(
     generator = torch.Generator().manual_seed(seed)
 
     id = "ID_0"
-    gender = gender_dict[which_id]
+    gender = gender_dict[selected_identity]
 
     # Construct prompt from dropdown selections
     prompt = f"face {pose.lower()} photo of {gender} sks person, {background.lower()} background"
@@ -87,38 +87,54 @@ with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(" # ID-Booth Demo")
 
+        with gr.Row():
+            gr.Markdown("### Choose an Identity:")
+
         with gr.Row():
             #gr.Markdown("### Choose an Identity:")
             identity_selectors = []
 
-            for id in id_list:
-                btn = gr.Image(
-                    value=os.path.join(folder_of_identity_images, id + ".jpg"),
-                    label=id,
-                    interactive=True,
-                    height=128,
-                    width=128,
-                )
-                identity_selectors.append(btn)
+            def make_identity_selector(identity, image_path):
+                with gr.Column():
+                    img = gr.Image(value=image_path, show_label=False, interactive=False, width=128, height=128)
+                    btn = gr.Button(value=f"Select {identity}")
+                    btn.click(
+                        fn=lambda: (identity, identity),
+                        inputs=[],
+                        outputs=[selected_identity],
+                    )
+
+            for identity in id_list:
+                image_path = os.path.join(folder_of_identity_images, identity + ".jpg")
+                make_identity_selector(identity, image_path)
+                
+            # for id in id_list:
+            #     btn = gr.Image(
+            #         value=os.path.join(folder_of_identity_images, id + ".jpg"),
+            #         label=id,
+            #         interactive=True,
+            #         height=128,
+            #         width=128,
+            #     )
+            #     identity_selectors.append(btn)
 
         # Set up click handlers
-        def select_identity(id):
-            return id
-
-        for btn, identity in zip(identity_selectors, folder_of_identity_images):
-            btn.select(
-                select_identity,
-                inputs=[],
-                outputs=[selected_identity],
-                _js=f"() => '{identity}'"
-            )
+        # def select_identity(id):
+        #     return id
+
+        # for btn, identity in zip(identity_selectors, folder_of_identity_images):
+        #     btn.select(
+        #         select_identity,
+        #         inputs=[],
+        #         outputs=[selected_identity],
+        #     )
             
         with gr.Row():
-            which_id = gr.Dropdown(
-                label="Identity",
-                choices=id_list,
-                value=id_list[0],
-            )
+            # which_id = gr.Dropdown(
+            #     label="Identity",
+            #     choices=id_list,
+            #     value=id_list[0],
+            # )
 
             background = gr.Dropdown(
                 label="Background",