Update app.py

app.py

@@ -4,24 +4,28 @@ from huggingface_hub import list_repo_files, hf_hub_url
 from collections import defaultdict
 import requests
 
-
 repo_id = 'nateraw/stable-diffusion-gallery'
 
-def get_data():
-    all_files = list_repo_files(repo_id, repo_type='dataset')
-    global data
+def get_data(revision='main'):
     data = defaultdict(list)
-    for file in all_files:
+    for file in list_repo_files(repo_id, repo_type='dataset', revision=revision):
         path = Path(file)
         if path.name == '.gitattributes':
             continue
         if path.suffix in ['.png', '.jpg', '.jpeg']:
             data[path.parent.name].append(file)
-    return gr.update(value=list(data.keys()))
-
-get_data()
-
-def fn(run):
+    print(data.keys())
+    return data
+
+def on_refresh(data_state):
+    data = get_data()
+    data_state.update(dict(data))
+    return [
+        data_state,
+        gr.update(choices=list(data_state.keys())),
+    ]
+
+def on_submit(run, data):
     images = [(hf_hub_url(repo_id=repo_id, filename=img, repo_type='dataset'), f"Seed:\n{Path(img).stem}") for img in data[run]]
     prompt_config_url = hf_hub_url(
         repo_id=repo_id,
@@ -31,13 +35,15 @@ def fn(run):
     prompt_config_json = requests.get(prompt_config_url).json()
     return prompt_config_json, images
 
+runs = list(get_data().keys())
 with gr.Blocks() as demo:
+    data_state = gr.State({})
     with gr.Column(variant="panel"):
         with gr.Row(variant="compact"):
             refresh_btn = gr.Button("Refresh").style(full_width=True)
 
         with gr.Row(variant="compact"):
-            run = gr.Dropdown(list(data.keys()), interactive=True).style(container=False)
+            run = gr.Dropdown(runs) #interactive=True) # .style(container=False)
             btn = gr.Button("View images").style(full_width=False)
 
         with gr.Row(variant="compact"):
@@ -46,7 +52,7 @@ with gr.Blocks() as demo:
             label="Generated images", show_label=False, elem_id="gallery"
         ).style(grid=[5], height="auto")
 
-    refresh_btn.click(get_data, None, run)
-    btn.click(fn, run, [data_json, gallery])
+    refresh_btn.click(on_refresh, data_state, [data_state, run])
+    btn.click(on_submit, [run, data_state], [data_json, gallery])
 
 demo.launch(debug=True)