Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
-from huggingface_hub import login
+from huggingface_hub import login, HfFileSystem
+
 from diffusers import DiffusionPipeline, StableDiffusionXLPipeline
 import torch
 import copy
@@ -11,8 +12,16 @@ is_shared_ui = True if "fffiloni/sd-xl-lora-fusion" in os.environ['SPACE_ID'] el
 hf_token = os.environ.get("HF_TOKEN")
 login(token = hf_token)
 
+fs = HfFileSystem(token=hf_token)
+
 original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
 
+def load_sfts(repo_1_id, repo_2_id):
+    # List all ".safetensors" files in repos
+    sfts_available_files_1 = fs.glob(f"{repo_1_id}/**.safetensors")
+    sfts_available_files_2 = fs.glob(f"{repo_2_id}/**.safetensors")
+    return gr.update(choices=sfts_available_files_1, value=sfts_available_files_1[0], visible=True), gr.update(choices=sfts_available_files_2, value=sfts_available_files_2[0], visible=True)
+    
 @spaces.GPU
 def infer(lora_1_id, lora_1_sfts, lora_2_id, lora_2_sfts, prompt, negative_prompt, lora_1_scale, lora_2_scale, seed):
 
@@ -141,16 +150,20 @@ with gr.Blocks(css=css) as demo:
             
             with gr.Column():
                 
-                lora_2_id = gr.Textbox(
+                lora_2_id = gr.Dropdown(
                     label = "LoRa 2 ID",
-                    placeholder = "username/model_id"
+                    placeholder = "username/model_id",
+                    visible=False
                 )
 
-                lora_2_sfts = gr.Textbox(
+                lora_2_sfts = gr.Dropdown(
                     label = "Safetensors file",
-                    placeholder = "specific_chosen.safetensors"
+                    placeholder = "specific_chosen.safetensors",
+                    visible=False
                 )
 
+        load_models_btn = gr.Button("Load models and .safetensors")
+
         # PART 2 • INFERENCE
         with gr.Row():
             
@@ -160,7 +173,7 @@ with gr.Blocks(css=css) as demo:
                 placeholde = "e.g: a triggerWordOne portrait in triggerWord2 style"
             )
 
-            run_btn = gr.Button("Run")
+        run_btn = gr.Button("Run")
         
         output_image = gr.Image(
             label = "Output"
@@ -205,6 +218,17 @@ with gr.Blocks(css=css) as demo:
             )
     
     # ACTIONS
+    load_models_btn.click(
+        fn = load_sfts, 
+        inputs = [
+            lora_1_id,
+            lora_2_id
+        ],
+        outputs = [
+            lora_1_sfts,
+            lora_2_sfts
+        ]
+    )
     run_btn.click(
         fn = infer,
         inputs = [