download results

app.py

@@ -54,6 +54,7 @@ with gr.Blocks(title="RFD3 Test") as demo:
     gen_directory = gr.State(None)
     gen_results = gr.State(None)
     gen_btn = gr.Button("Run Unconditional Generation")
+    output_file = gr.File(label="Download RFD3 results as zip", visible=True)
 
 
     # Section to inspect PDB of generated structures
@@ -73,11 +74,21 @@ with gr.Blocks(title="RFD3 Test") as demo:
     display_state = gr.Textbox(label="Selected Batch and Design", visible=True)
     display_state.value = "Please Select a Batch and Design number to show sequence"
 
+    def download_results_as_zip(directory):
+        if directory is None:
+            return gr.update()
+        zip_path = f"{directory}.zip"
+        shutil.make_archive(directory, 'zip', directory)
+        return gr.update(value=zip_path, visible=True)
         
     gen_btn.click(unconditional_generation, inputs=[num_batches, num_designs_per_batch, length], outputs=[gen_directory, gen_results]).then(
         update_batch_choices,
         inputs=gen_results,
-        outputs=batch_dropdown)
+        outputs=batch_dropdown).then(
+            download_results_as_zip,
+            inputs=gen_directory,
+            outputs=output_file
+        )
 
     batch_dropdown.change(update_designs, inputs=[batch_dropdown, gen_results], outputs=[design_dropdown])
     design_dropdown.change()
@@ -85,6 +96,9 @@ with gr.Blocks(title="RFD3 Test") as demo:
     show_pdb_btn.click(show_pdb, inputs=[batch_dropdown, design_dropdown, gen_results], outputs=display_state)
 
 
+
+
+
     #def load_viewer(batch, design, result):
     #    if batch is None or design is None:
     #        return gr.update()