Fix gallery behavior when adding custom LoRA (#3)

- Fix gallery behavior when adding custom LoRA (14408dadee3ffe7a77cb31966c4c8b67b20bbc93)


Co-authored-by: John Smith <[email protected]>

app.py

@@ -170,7 +170,7 @@ def randomize_loras(selected_indices, loras_state):
     random_prompt = random.choice(prompt_values)
     return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, random_prompt
 
-def add_custom_lora(custom_lora, selected_indices, current_loras):
+def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
     if custom_lora:
         try:
             title, repo, path, trigger_word, image = check_custom_model(custom_lora)
@@ -232,7 +232,7 @@ def add_custom_lora(custom_lora, selected_indices, current_loras):
     else:
         return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
 
-def remove_custom_lora(selected_indices, current_loras):
+def remove_custom_lora(selected_indices, current_loras, gallery):
     if current_loras:
         custom_lora_repo = current_loras[-1]['repo']
         # Remove from loras list
@@ -562,12 +562,12 @@ with gr.Blocks(css=css, delete_cache=(60, 3600)) as app:
     )
     add_custom_lora_button.click(
         add_custom_lora,
-        inputs=[custom_lora, selected_indices, loras_state],
+        inputs=[custom_lora, selected_indices, loras_state, gallery],
         outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
     remove_custom_lora_button.click(
         remove_custom_lora,
-        inputs=[selected_indices, loras_state],
+        inputs=[selected_indices, loras_state, gallery],
         outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
     gr.on(