integrate download zip logic inside generation_with_input_config, remove gr.update

app.py

@@ -64,7 +64,7 @@ with gr.Blocks(title="RFD3 Test") as demo:
 
                 """)
 
-    gen_directory = gr.State(None) # the directory where generation results are saved, used to trigger the download of results as zip file
+
     gen_results = gr.State(None) # the results of the generation, which is a list of dicts where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path".
 
     # inputs from user
@@ -130,23 +130,17 @@ 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"
-    
-
         
     #generation_event = run_btn.click(
     #    lambda: (gr.update(visible=False), gr.update(visible=True)),
     #    outputs=[run_btn, stop_btn]
     #).then(
     generation_event = run_btn.click(
-        generation_with_input_config, inputs=[config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args, max_duration], outputs=[runtextbox, gen_directory, gen_results]
+        generation_with_input_config, inputs=[config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args, max_duration], outputs=[runtextbox, gen_results, output_file]
     ).then(
         update_batch_choices,
         inputs=gen_results,
         outputs=batch_dropdown
-    ).then(
-        download_results_as_zip,
-        inputs=gen_directory,
-        outputs=output_file
     )#.then(
     #    lambda: (gr.update(visible=True), gr.update(visible=False)),
     #    outputs=[run_btn, stop_btn]

utils/handle_files.py

@@ -33,18 +33,18 @@ def download_results_as_zip(directory):
     
     Parameters:
     ----------
-    directory: gr.State
+    directory: gr.State or str
         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 or None: Path to the created zip file if directory is valid, else None.
     """
     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
 
 
 def collect_outputs(gen_directory, num_batches, num_designs_per_batch):

utils/pipelines.py

@@ -42,7 +42,7 @@ def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_
 
 
     if input_file is None:        
-        return gr.update(value="Please ensure you have uploaded a configuration file: .yaml or .json"), gr.update(), gr.update()
+        return "Please ensure you have uploaded a configuration file: .yaml or .json", None, None
     elif pdb_file is None:
         status_update = f"Running generation for {num_batches} batches of {num_designs_per_batch}\n job configuration uploaded from file {os.path.basename(input_file)}\n no scaffold/target provided"
     else:
@@ -67,11 +67,9 @@ def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_
             command = f"rfd3 design inputs={input_file} out_dir={directory} n_batches={num_batches} diffusion_batch_size={num_designs_per_batch}"
         if extra_args:
             command += f" {extra_args}"
-        print(f"Running command: {command}")
         status_update += f"\nRunning command: {command}."
         start = perf_counter()
         res = subprocess.run(command, shell=True, check=True, text=True, capture_output=True)
-        print("Command took", perf_counter() - start, "seconds to run.")
         status_update += f"\nGeneration successful! Command took {perf_counter() - start:.2f} seconds to run."
 
 
@@ -86,13 +84,13 @@ def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_
                 cif_path = os.path.join(directory, file_name)
                 pdb_path = mcif_gz_to_pdb(cif_path)
                 results.append({"batch": batch, "design": design, "cif_path": cif_path, "pdb_path": pdb_path})
-            
-        print(results)
-        return gr.update(value=status_update), directory, results
+        
+        zip_path = download_results_as_zip(directory)
+
+        return status_update, results, zip_path
     
     except subprocess.CalledProcessError as e:
-        print("subprocess threw an error", e.stderr)
-        return gr.update(value=f"Generation failed:\n{e.stderr}"), None, None
+        return f"Generation failed:\n{e.stderr}", None, None
 
 
 #def generation_with_input_config_factory(max_duration):