Upload example images

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import numpy as np
 import random
 import spaces  # Uncomment if using ZeroGPU
-
+import os 
 from diffusers import StableDiffusionPipeline, DDPMScheduler
 import torch
 
@@ -21,12 +21,14 @@ folder_of_lora_weights = "./ID-Booth_LoRA_weights"
 which_checkpoint = "checkpoint-31-6400"
 lora_name = "pytorch_lora_weights.safetensors"
 
+selected_identity = gr.State(value="ID_0")  # Default selection
+folder_of_identity_images = "./assets/example_images/"
 
 backgrounds_list = ["forest", "city street", "beach", "office", "bus", "laboratory", "factory", "construction site", "hospital", "night club", ""]
 poses_list = ["portrait", "side-portrait"]
-id_list = ["ID_0", "ID_1", "ID_2"]
+id_list = ["ID_0", "ID_1", "ID_2", "ID_3"]
 
-gender_dict = {"ID_0": "male", "ID_1": "male", "ID_2": "female", "ID_2": "male"}
+gender_dict = {"ID_1": "male", "ID_5": "male", "ID_16": "female", "ID_20": "male"}
 MAX_SEED = 10000
 image_size = 512
 
@@ -85,6 +87,32 @@ 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:")
+            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)
+
+        # 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}'"
+            )
+            
         with gr.Row():
             which_id = gr.Dropdown(
                 label="Identity",