fix how number of abtches and designs is passed to rfd3 CLI

app.py

@@ -53,6 +53,20 @@ with gr.Blocks(title="RFD3 Test") as demo:
             with gr.Tabs() as config_tabs:
                 with gr.TabItem("Upload Config") as upload_tab: # upload a config yaml or json
                     config_upload = gr.File(label="Config file: .yaml or .json", file_types=[".pdb", ".yaml", ".json"])
+                    num_designs_per_batch_upload = gr.Number(
+                        value=2,
+                        label="Number of Designs per Batch",
+                        precision=0,
+                        minimum=1,
+                        maximum=8
+                    )
+                    num_batches_upload = gr.Number(
+                        value=5,
+                        label="Number of Batches",
+                        precision=0,
+                        minimum=1,
+                        maximum=10
+                    )
                 with gr.TabItem("Manual Config") as manual_tab: # minimal config for testing unconditional generation
                     num_designs_per_batch = gr.Number(
                         value=2,
@@ -96,7 +110,7 @@ with gr.Blocks(title="RFD3 Test") as demo:
     manual_tab.select(lambda: "manual", outputs=tab_selected)
 
     # validate the configuration
-    config_validation_btn.click(validate_config_ready, inputs=[tab_selected, config_upload, num_designs_per_batch, num_batches, length], outputs=[config_textbox, config_ready])
+    config_validation_btn.click(validate_config_ready, inputs=[tab_selected, config_upload, num_designs_per_batch, num_batches, length, num_designs_per_batch_upload, num_batches_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])
 
@@ -120,18 +134,18 @@ 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 generate(config_ready, scaffold_ready, num_batches, num_designs_per_batch, length, config_upload):
+    def generate(config_ready, scaffold_ready, num_batches, num_designs_per_batch, length, config_upload, num_designs_per_batch_upload, num_batches_upload):
         if config_ready is None or scaffold_ready is None:
             return None, None
         if config_ready == "upload" and scaffold_ready == "no_input":
-            gen_directory, gen_results = unconditional_generation_with_input_config(config_upload)
+            gen_directory, gen_results = unconditional_generation_with_input_config(config_upload, num_batches_upload, num_designs_per_batch_upload)
             return gen_directory, gen_results
         if config_ready=="manual" and scaffold_ready=="no_input":
             gen_directory, gen_results = unconditional_generation(num_batches, num_designs_per_batch, length)
             return gen_directory, gen_results
         
     run_btn.click(give_run_status, inputs=[config_ready, scaffold_ready, num_batches, num_designs_per_batch, length, config_upload], outputs=runtextbox).then(
-        generate, inputs=[config_ready, scaffold_ready, num_batches, num_designs_per_batch, length, config_upload], outputs=[gen_directory, gen_results]
+        generate, inputs=[config_ready, scaffold_ready, num_batches, num_designs_per_batch, length, config_upload, num_designs_per_batch_upload, num_batches_upload], outputs=[gen_directory, gen_results]
     ).then(
         update_batch_choices,
         inputs=gen_results,

utils/handle_events.py

@@ -58,7 +58,7 @@ 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(tab_selected, config_upload, num_designs_per_batch, num_batches, length):
+def validate_config_ready(tab_selected, config_upload, num_designs_per_batch, num_batches, length, num_designs_per_batch_upload, num_batches_upload):
     """
     Once the user presses on the Validate Config button, this function checks whether the config is really ready to go.
     It is ready, either if the user has selected the manual config tab, or if they have uploaded a config file in the upload config tab and have it open.
@@ -86,7 +86,10 @@ def validate_config_ready(tab_selected, config_upload, num_designs_per_batch, nu
     elif tab_selected == "upload" and config_upload is None:
         return "Please upload a config file or switch to manual config for uncondition generation.", None
     elif tab_selected == "upload" and config_upload is not None:
-        return "Config file uploaded and valid!", "upload"
+        for num in [num_designs_per_batch_upload, num_batches_upload]:
+            if num is None:
+                return "Please specify the number of batches and designs per batch for your configuration based on the uploaded config file", None
+        return "Config file uploaded and valid! Number of batches and designs validated!", "upload"
     
 def validate_scaffold_ready_with_file(scaffold_upload):
     """

utils/pipelines.py

@@ -99,7 +99,7 @@ def unconditional_generation(num_batches, num_designs_per_batch, length):
         raise RuntimeError(f"Error during generation: {str(e)}")
     
 @spaces.GPU(duration=300) 
-def unconditional_generation_with_input_config(input_file):
+def unconditional_generation_with_input_config(input_file, num_batches, num_designs_per_batch):
     """
     Runs an unconditional generation with the specified input config file. Saves the generated structures to a timestamped directory in the outputs folder and returns the path to the directory along with a list of the generated structures' file paths.
 
@@ -124,7 +124,7 @@ def unconditional_generation_with_input_config(input_file):
 
     try: 
         
-        command = f"rfd3 design inputs={input_file} out_dir={directory}"
+        command = f"rfd3 design inputs={input_file} out_dir={directory} n_batches={num_batches} diffusion_batch_size={num_designs_per_batch}"
         print(f"Running command: {command}")
         subprocess.run(command, shell=True, check=True, capture_output=True, text=True)