simplify status update string by integrating in generate, create stop button

app.py

@@ -64,42 +64,31 @@ with gr.Blocks(title="RFD3 Test") as demo:
 
                 """)
 
-    #config_ready = gr.State(None) # whether the config is ready for running generation, takes values None, "manual", or "upload"
-    #scaffold_ready = gr.State(None) # whether the scaffold is ready for running generation, takes values None, "upload", or "no_input"
     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
     with gr.Row():
         with gr.Column(scale=1):  # Left half
-            #gr.Markdown("Set up the configuration for your run through a valid yaml file. Choose the number of batches and designs per batch for your run.")
             config_upload = gr.File(label="Config file: .yaml or .json", file_types=[".pdb", ".yaml", ".json"])
-            #config_validation_btn = gr.Button("Validate Config")
-            #config_textbox = gr.Textbox(label="Configuration status", value ="Waiting for config validation...")
-
 
         with gr.Column(scale=1):  # Right half
-            #gr.Markdown("Upload your target/scaffold structure as a PDB file to condition the generation. Press 'No Scaffold/Target' if you want to run an unconditional generation.")
             scaffold_upload = gr.File(label="Target/Scaffold PDB", file_types=[".pdb"])
-            #with gr.Row():
-            #    scaffold_validation_btn = gr.Button("Validate Scaffold")
-            #    no_input_btn = gr.Button("No Scaffold/Target")
-            #scaffold_textbox = gr.Textbox(label="Scaffold/Target status", value ="Waiting for scaffold validation...")
 
     with gr.Row():
             num_designs_per_batch = gr.Number(
-                value=2,
+                value=8,
                 label="Number of Designs per Batch",
                 precision=0,
                 minimum=1,
                 maximum=16
             )
             num_batches= gr.Number(
-                value=5,
+                value=2,
                 label="Number of Batches",
                 precision=0,
                 minimum=1,
-                maximum=10
+                maximum=50
             )
             max_duration = gr.Number(
                 value=3600,
@@ -117,15 +106,12 @@ with gr.Blocks(title="RFD3 Test") as demo:
             info="Add extra RFD3 CLI arguments here (optional)"
         )
 
-    run_btn = gr.Button("Run Generation", variant="primary")
+    with gr.Row():
+        run_btn = gr.Button("Run Generation", variant="primary")
+        stop_btn = gr.Button("Stop Generation", variant="stop", visible=False)
     runtextbox = gr.Textbox(label="Run status", value="Waiting for generation run...")
 
 
-    # validate the configuration
-    #config_validation_btn.click(validate_config_ready, inputs=config_upload, outputs=[config_textbox, config_ready])
-    #scaffold_validation_btn.click(validate_scaffold_ready_with_file, inputs=scaffold_upload, outputs=[scaffold_textbox, scaffold_ready])
-    #no_input_btn.click(lambda: ("No scaffold/target will be used. Ready for unconditional generation!", "no_input"), outputs=[scaffold_textbox, scaffold_ready])
-
     output_file = gr.File(label="Download RFD3 results as zip", visible=True)
 
     # Section to inspect PDB of generated structures
@@ -149,12 +135,17 @@ with gr.Blocks(title="RFD3 Test") as demo:
         if config_upload is None:
             return gr.update(), None, None
         else:
-            #generate_with_correct_time_limit = generation_with_input_config_factory(max_duration)
             textbox_update, gen_directory, gen_results = generation_with_input_config(config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args, max_duration)
             print(textbox_update)
             return textbox_update, gen_directory, gen_results
         
-    run_btn.click(give_run_status, inputs=[num_batches, num_designs_per_batch, config_upload, scaffold_upload], outputs=runtextbox).then(
+    generation_event = run_btn.click(
+        lambda: gr.update(visible=False),
+        outputs=run_btn
+    ).then(
+        lambda: gr.update(visible=True),
+        outputs=stop_btn
+    ).then(
         generate, inputs=[config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args, max_duration], outputs=[runtextbox, gen_directory, gen_results]
     ).then(
         update_batch_choices,
@@ -165,9 +156,23 @@ with gr.Blocks(title="RFD3 Test") as demo:
         inputs=gen_directory,
         outputs=output_file
     ).then(
-        lambda gen_results: gr.update(value="Generation completed!") if gen_results is not None else gr.update(), 
-        inputs=gen_results,
-        outputs=runtextbox
+        lambda: gr.update(visible=True),
+        outputs=run_btn
+    ).then(
+        lambda: gr.update(visible=False),
+        outputs=stop_btn
+    )
+    
+    stop_btn.click(
+        lambda: gr.update(value="Generation cancelled by user."),
+        outputs=runtextbox,
+        cancels=[generation_event]
+    ).then(
+        lambda: gr.update(visible=True),
+        outputs=run_btn
+    ).then(
+        lambda: gr.update(visible=False),
+        outputs=stop_btn
     )
         
 

utils/handle_events.py

@@ -59,71 +59,72 @@ def show_pdb(batch, design, result):
     print(pdb_str)
     return gr.update(value=f"Selected Batch: {batch}, Design: {design}, saved at {pdb_path}:\n {pdb_str}", visible=True)
 
-def validate_config_ready(config_upload):
-    """
-    Once the user presses on the Validate Config button, this function checks whether the config is really ready to go.
-    It is ready if the config file is uploaded.
-
-    Parameters:
-    ----------
-
-    config_upload: gr.File
-        uploaded config file, is None if no file is uploaded
-
-
-    Returns:
-    -------
-    str: message to the user about the status of the config validation
-    str: "upload" if upload, and None if config is not ready for generation.
-    """
-    if config_upload is None:
-        return "Please upload a config file for the generation to be ready.", None
-    return "Config file uploaded!", "upload"
-
-    
-def validate_scaffold_ready_with_file(scaffold_upload):
-    """
-    Once the user presses on the Validate Scaffold button, this function checks whether a file was indeed uploaded. If not, it prompts the user to upload a file or select no scaffold/target for unconditioned generation.
-
-    Parameters:
-    ----------
-    scaffold_upload: gr.File
-        uploaded scaffold file, is None if no file is uploaded
-
-    Returns:
-    -------
-    str: message to the user about the status of the scaffold validation
-    str: status string. "upload" if scaffold file is uploaded and valid, None if not.
-    """
-    if scaffold_upload is not None:
-        return "Scaffold file uploaded and valid!", "upload"
-    else:
-        return "Please upload a scaffold file or press 'No Scaffold/Target' for unconditional generation.", None
-    
-def give_run_status(num_batches, num_designs_per_batch, config_upload, scaffold_upload):
-    """
-    Once the user presses on the Run Generation button, this function checks whether both config and scaffold are ready, and gives a status message about the generation run that is about to happen. If both config and scaffold are not ready, it prompts the user to ensure they are both validated before running the generation.
-
-    Parameters:
-    ----------
-    num_batches: gr.Number
-        number of batches input by the user manually
-    num_designs_per_batch: gr.Number
-        number of designs per batch input by the user manually
-    config_upload: gr.File
-        uploaded config file, is None if no file is uploaded
-    scaffold_upload: gr.File
-        uploaded scaffold file, is None if no file is uploaded
-
-    Returns:
-    -------
-        gr.update: Gradio update object to update the generation status message in the UI about to be run. If both config and scaffold are not ready, returns an update with a message prompting the user to ensure they are both validated before running the generation.
-
-    
-    """
-    if config_upload is None:        
-        return gr.update(value="Please ensure you have uploaded a configuration file: .yaml or .json")
-    elif scaffold_upload is None:
-        return gr.update(value=f"Running unconditional generation for {num_batches} batches of {num_designs_per_batch}, config uploaded from file {os.path.basename(config_upload)}")
-    else:
-        return gr.update(value=f"Running conditional generation for {num_batches} batches of {num_designs_per_batch}, config uploaded from file {os.path.basename(config_upload)}, scaffold uploaded as {os.path.basename(scaffold_upload)}")
+#def validate_config_ready(config_upload):
+#    """
+#    Once the user presses on the Validate Config button, this function checks whether the config is really ready to go.
+#    It is ready if the config file is uploaded.
+#
+#    Parameters:
+#    ----------
+#
+#    config_upload: gr.File
+#        uploaded config file, is None if no file is uploaded
+#
+#
+#    Returns:
+#    -------
+#    str: message to the user about the status of the config validation
+#    str: "upload" if upload, and None if config is not ready for generation.
+#    """
+#    if config_upload is None:
+#        return "Please upload a config file for the generation to be ready.", None
+#    return "Config file uploaded!", "upload"
+#
+#    
+#def validate_scaffold_ready_with_file(scaffold_upload):
+#    """
+#    Once the user presses on the Validate Scaffold button, this function checks whether a file was indeed uploaded. If not, it prompts the user to upload a file or select no scaffold/target for unconditioned generation.
+#
+#    Parameters:
+#    ----------
+#    scaffold_upload: gr.File
+#        uploaded scaffold file, is None if no file is uploaded
+#
+#    Returns:
+#    -------
+#    str: message to the user about the status of the scaffold validation
+#    str: status string. "upload" if scaffold file is uploaded and valid, None if not.
+#    """
+#    if scaffold_upload is not None:
+#        return "Scaffold file uploaded and valid!", "upload"
+#    else:
+#        return "Please upload a scaffold file or press 'No Scaffold/Target' for unconditional generation.", None
+#
+#
+#def give_run_status(num_batches, num_designs_per_batch, config_upload, scaffold_upload):
+#    """
+#    Once the user presses on the Run Generation button, this function checks whether both config and scaffold are ready, and gives a status message about the generation run that is about to happen. If both config and scaffold are not ready, it prompts the user to ensure they are both validated before running the generation.
+#
+#    Parameters:
+#    ----------
+#    num_batches: gr.Number
+#        number of batches input by the user manually
+#    num_designs_per_batch: gr.Number
+#        number of designs per batch input by the user manually
+#    config_upload: gr.File
+#        uploaded config file, is None if no file is uploaded
+#    scaffold_upload: gr.File
+#        uploaded scaffold file, is None if no file is uploaded
+#
+#    Returns:
+#    -------
+#        gr.update: Gradio update object to update the generation status message in the UI about to be run. If both config and scaffold are not ready, returns an update with a message prompting the user to ensure they are both validated before running the generation.
+#
+#    
+#    """
+#    if config_upload is None:        
+#        return gr.update(value="Please ensure you have uploaded a configuration file: .yaml or .json")
+#    elif scaffold_upload is None:
+#        return gr.update(value=f"Running unconditional generation for {num_batches} batches of {num_designs_per_batch}, config uploaded from file {os.path.basename(config_upload)}")
+#    else:
+#        return gr.update(value=f"Running conditional generation for {num_batches} batches of {num_designs_per_batch}, config uploaded from file {os.path.basename(config_upload)}, scaffold uploaded as {os.path.basename(scaffold_upload)}")

utils/pipelines.py

@@ -40,6 +40,14 @@ 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()
+    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:
+        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 scaffold/target provided from file {os.path.basename(pdb_file)}"
+
     session_hash = random.getrandbits(128)
     time_stamp = time.strftime("%Y-%m-%d-%H-%M-%S")
     directory = f"./outputs/generation_with_input_config/session_{session_hash}_{time_stamp}"
@@ -60,9 +68,11 @@ def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_
         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."
 
 
         results = []
@@ -78,7 +88,7 @@ def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_
                 results.append({"batch": batch, "design": design, "cif_path": cif_path, "pdb_path": pdb_path})
             
         print(results)
-        return gr.update(value="Generation Successful"), directory, results
+        return gr.update(value=status_update), directory, results
     
     except subprocess.CalledProcessError as e:
         print("subprocess threw an error", e.stderr)