Update app.py

app.py

@@ -14,18 +14,26 @@ def get_data(revision='main'):
             continue
         if path.suffix in ['.png', '.jpg', '.jpeg']:
             data[path.parent.name].append(file)
-    print(data.keys())
-    return data
+
+    prompts = []
+    for key in sorted(data.keys()):
+        prompt_cfg_url = hf_hub_url(repo_id=repo_id, filename=f"{key}/prompt_config.json", repo_type='dataset')
+        prompt_cfg = requests.get(prompt_cfg_url).json()
+        text = prompt_cfg.get('prompt')
+        prompts.append(text)
+
+    return data, prompts
 
 def on_refresh(data_state):
-    data = get_data()
+    data, prompts = get_data()
     data_state.update(dict(data))
     return [
         data_state,
-        gr.update(choices=sorted(data_state.keys())),
+        gr.update(choices=prompts),
     ]
 
-def on_submit(run, data):
+def on_submit(run_idx, data):
+    run = sorted(data.keys())[run_idx]
     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,
@@ -35,15 +43,16 @@ def on_submit(run, data):
     prompt_config_json = requests.get(prompt_config_url).json()
     return prompt_config_json, images
 
-runs = list(get_data().keys())
+data, prompts = get_data()
+runs = sorted(data.keys())
 with gr.Blocks() as demo:
-    data_state = gr.State({})
+    data_state = gr.State(data)
     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(runs) #interactive=True) # .style(container=False)
+            run_idx = gr.Dropdown(prompts, type='index') #interactive=True) # .style(container=False)
             btn = gr.Button("View images").style(full_width=False)
 
         with gr.Row(variant="compact"):
@@ -52,7 +61,8 @@ with gr.Blocks() as demo:
             label="Generated images", show_label=False, elem_id="gallery"
         ).style(grid=[5], height="auto")
 
-    refresh_btn.click(on_refresh, data_state, [data_state, run])
-    btn.click(on_submit, [run, data_state], [data_json, gallery])
+    refresh_btn.click(on_refresh, data_state, [data_state, run_idx])
+    btn.click(on_submit, [run_idx, data_state], [data_json, gallery])
 
-demo.launch(debug=True)
+if __name__ == '__main__':
+    demo.launch(debug=True)