Update UI

app.py

@@ -21,7 +21,6 @@ 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", ""]
@@ -34,7 +33,7 @@ image_size = 512
 
 @spaces.GPU  # Uncomment if using ZeroGPU
 def infer(
-    selected_identity,
+    identity,
     background,
     pose,
     negative_prompt,
@@ -46,7 +45,7 @@ def infer(
     num_images=1
 ):
     
-    full_lora_weights_path = f"{folder_of_lora_weights}/{selected_identity}/{which_checkpoint}/{lora_name}"
+    full_lora_weights_path = f"{folder_of_lora_weights}/{identity}/{which_checkpoint}/{lora_name}"
     pipe.load_lora_weights(full_lora_weights_path)
 
     if randomize_seed:
@@ -54,8 +53,7 @@ def infer(
 
     generator = torch.Generator().manual_seed(seed)
 
-    id = "ID_0"
-    gender = gender_dict[selected_identity]
+    gender = gender_dict[identity]
 
     # Construct prompt from dropdown selections
     prompt = f"face {pose.lower()} photo of {gender} sks person, {background.lower()} background"