include zip_file download logic on fold_all_jobs, remove gr.update returns

app.py

@@ -62,8 +62,6 @@ with gr.Blocks(title="RFD3 Test") as demo:
                 all taken from <a href="https://github.com/RosettaCommons/foundry/tree/production/models/rf3" target="_blank">RosettaCommons' foundry repository.</a> 
                 """)
 
-    gen_directory = gr.State(None) # the directory where generation results are saved, used to trigger the download of results as zip file
-
     # inputs from user
     with gr.Row():
         with gr.Column(scale=1):  # Left half
@@ -109,16 +107,10 @@ with gr.Blocks(title="RFD3 Test") as demo:
     runtextbox = gr.Textbox(label="Run status", value="Waiting for folding run...")
     output_file = gr.File(label="Download RF3 results as zip", visible=True)
 
-
         
     run_btn.click(
-        fold_all_jobs, inputs=[job_upload, support_upload, num_predictions, early_stopping, diffusion_steps, max_duration], outputs=[runtextbox, gen_directory]
-    ).then(
-        download_results_as_zip,
-        inputs=gen_directory,
-        outputs=output_file
+        fold_all_jobs, inputs=[job_upload, support_upload, num_predictions, early_stopping, diffusion_steps, max_duration], outputs=[runtextbox, output_file]
     )
-        
 
 
 if __name__ == "__main__":

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: path to the generated zip file, to be passed to gr.File for download.
     """
     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)
+    return zip_path
 
     
 

utils/pipelines.py

@@ -32,17 +32,19 @@ def fold_all_jobs(job_files, support_files, num_predictions, early_stopping, dif
         The pLDDT threshold for early stopping. If set to 0, early stopping is disabled.
     diffusion_steps: int
         The number of diffusion steps to use during generation. Higher values may improve structure quality but will increase runtime.
+    max_duration: int
+        The maximum duration (in seconds) for the folding process. This is used to set the timeout for the Spaces GPU allocation on ZeroGPU.
     
     Returns:
     -------
         textbox_update: gr.update,
             A gr.update object to update the textbox with the status of the generation. propagates subprocess errors to the textbox if the generation fails.
-        directory: str,
-            The path to the directory where the generated structures are saved.
+        zip_path: str or None
+            If generation is successful, returns the path to the generated zip file containing the results, which can be passed to gr.File for download. None if generation fails.
     """
 
     if job_files is None:        
-        return gr.update(value="Please ensure you have uploaded at least one job file: .json, .pdb, or .cif"), None
+        return "Please ensure you have uploaded at least one job file: .json, .pdb, or .cif", None
     elif support_files is None:
         status_update = f"Creating {num_predictions} structure predictions per job.\n jobs uploaded {[os.path.basename(f.name) for f in job_files]}\nno support files uploaded."
     else:
@@ -68,8 +70,10 @@ def fold_all_jobs(job_files, support_files, num_predictions, early_stopping, dif
         print("Command took", perf_counter() - start, "seconds to run.")
         status_update += f"\nGeneration successful! Command took {perf_counter() - start:.2f} seconds to run."
 
-        return gr.update(value=status_update), directory
+        zip_path = download_results_as_zip(directory)
+
+        return status_update, 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
+        return f"Generation failed:\n{e.stderr}", None