Update app.py

app.py

@@ -16,10 +16,27 @@ fs = HfFileSystem(token=hf_token)
 
 original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
 
+def get_files(file_paths):
+    last_files = {}  # Dictionary to store the last file for each path
+
+    for file_path in file_paths:
+        # Split the file path into directory and file components
+        directory, file_name = file_path.rsplit('/', 1)
+    
+        # Update the last file for the current path
+        last_files[directory] = file_name
+    
+    # Extract the last files from the dictionary
+    result = list(last_files.values())
+
+    return result
+
 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")
+    sfts_available_files_1 = fs.ls(f"{repo_1_id}/tree/main**.safetensors", detail=False)
+    sfts_available_files_1 = get_files(sfts_available_files_1)
+    sfts_available_files_2 = fs.ls(f"{repo_2_id}/tree/main**.safetensors", detail=False)
+    sfts_available_files_2 = get_files(sfts_available_files_2)
     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