integrate download_zip logics inside run_generation_folder

app.py

@@ -107,11 +107,9 @@ with gr.Blocks(title="") as demo:
             gen_directory_multi = gr.State(value=None)
             
             run_btn_multi.click(
-                run_generation_folder, inputs=[pdb_folder, num_batches_multi, num_designs_per_batch_multi, chains_to_design_multi, temperature_multi, extra_args_multi, max_duration], outputs=[runtextbox_multi, gen_directory_multi]
+                run_generation_folder, inputs=[pdb_folder, num_batches_multi, num_designs_per_batch_multi, chains_to_design_multi, temperature_multi, extra_args_multi, max_duration], outputs=[runtextbox_multi, gen_directory_multi, output_file_multi]
             ).then(
                 display_fasta, inputs=gen_directory_multi, outputs=fastatextbox_multi
-            ).then(
-                download_results_as_zip, inputs=gen_directory_multi, outputs=output_file_multi
             ).then(
                 lambda gen_dir: gr.update(choices=os.listdir(os.path.join(gen_dir, "seqs"))) if gen_dir else gr.update(), inputs=gen_directory_multi, outputs=select_fasta_to_show
             )

space_utils/handle_files.py

@@ -39,15 +39,15 @@ def download_results_as_zip(directory):
     
     Parameters:
     ----------
-    directory: gr.State
+    directory: str or gr.Path
         Path to the directory containing generated results. None if generation has not been run yet.
     
     Returns
     -------
-        gr.update: Gradio update object to trigger file download. If directory is None, returns an empty update.
+        str: path to the created zip file for download, or None if no directory was provided.
     """
     if directory is None:
-        return gr.update()
+        return None
     zip_path = f"{directory}.zip"
     shutil.make_archive(directory, 'zip', directory)
-    return gr.update(value=zip_path, visible=True)
+    return zip_path

space_utils/pipelines.py

@@ -6,7 +6,7 @@ import gradio as gr
 import shutil
 import json
 import spaces
-
+from space_utils.handle_files import display_fasta, download_results_as_zip
 #def run_generation_single_pdb(pdb_file, num_batches, num_designs_per_batch, chains_to_design, temperature, extra_args):
 #
 #    random_hash = random.getrandbits(128)
@@ -34,7 +34,7 @@ def get_duration(pdb_folder, num_batches, num_designs_per_batch, chains_to_desig
 def run_generation_folder(pdb_folder, num_batches, num_designs_per_batch, chains_to_design, temperature, extra_args, max_duration):
 
     if pdb_folder is None:
-        return "Please upload a folder of PDB files to run generation.", None
+        return "Please upload a folder of PDB files to run generation.", None, None
     
     status_string = f"Running generation with batch size {num_designs_per_batch} and number of batches {num_batches} for {len(pdb_folder)} pdb files..."
 
@@ -77,15 +77,9 @@ def run_generation_folder(pdb_folder, num_batches, num_designs_per_batch, chains
         res = subprocess.run(command, shell=True, check=True, text=True, capture_output=True)
         status_string += "\nGeneration complete!"
 
-        return status_string, out_dir
+        zip_path = download_results_as_zip(out_dir)
+
+        return status_string, out_dir, zip_path
     except subprocess.CalledProcessError as e:
-        return status_string+ f"\nGeneration failed: \n{e.stderr}", None
-    
-#def ligandmpnn_endpoint(pdb_folder, num_batches, num_designs_per_batch, chains_to_design, temperature, extra_args, max_duration):
-#    output_dir, status_update = gr.State(value=None), gr.State(value ="Waiting for generation run...")
-#    output_file_multi = gr.File(label="Download")
-#    output_dir, status_update =  
+        return status_string+ f"\nGeneration failed: \n{e.stderr}", None, None
 
-def ligandmpnn_endpoint(pdb_folder, num_batches, num_designs_per_batch, chains_to_design, temperature, extra_args, max_duration):
-    status_str, output_dir = run_generation_folder(pdb_folder, num_batches, num_designs_per_batch, chains_to_design, temperature, extra_args, max_duration)
-    return output_dir