try dynamically building the correct decorator

app.py

@@ -9,7 +9,7 @@ from atomworks.io.utils.visualize import view
 from lightning.fabric import seed_everything
 from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine
 from utils.download_weights import download_weights
-from utils.pipelines import test_rfd3_from_notebook, unconditional_generation, generation_with_input_config
+from utils.pipelines import *
 #from gradio_molecule3d import Molecule3D
 from utils.handle_events import *
 from utils.handle_files import *
@@ -101,6 +101,13 @@ with gr.Blocks(title="RFD3 Test") as demo:
                 minimum=1,
                 maximum=10
             )
+            max_duration = gr.Number(
+                value=240,
+                label="Max Duration (seconds)",
+                precision=0,
+                minimum=1,
+                maximum=7199
+            )
 
     with gr.Accordion(label="Advanced Options", open=False):
         extra_args = gr.Textbox(
@@ -138,16 +145,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"
     
-    def generate(config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args):
+    def generate(config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args, max_duration):
         if config_upload is None:
             return gr.update(), None, None
         else:
-            textbox_update, gen_directory, gen_results = generation_with_input_config(config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args)
+            generate_with_correct_time_limit = generation_with_input_config_factory(max_duration)
+            textbox_update, gen_directory, gen_results = generate_with_correct_time_limit(config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args)
             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(
-        generate, inputs=[config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args], outputs=[runtextbox, gen_directory, gen_results]
+        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,
         inputs=gen_results,

utils/pipelines.py

@@ -12,97 +12,10 @@ import shutil
 from time import perf_counter
 import random
 
-@spaces.GPU(duration=240)  
-def test_rfd3_from_notebook():
-    # Set seed for reproducibility
-    seed_everything(0)
 
-    # Configure RFD3 inference
-    config = RFD3InferenceConfig(
-        specification={
-            'length': 40,  # Generate 80-residue proteins
-        },
-        diffusion_batch_size=2,  # Generate 2 structures per batch
-    )
-
-    # Initialize engine and run generation
-    try: 
-        model = RFD3InferenceEngine(**config)
-        outputs = model.run(
-            inputs=None,      # None for unconditional generation
-            out_dir=None,     # None to return in memory (no file output)
-            n_batches=1,      # Generate 1 batch
-        )
-        return_str = "RDF3 test passed! Generated structures:\n"
-
-        for idx, data in outputs.items():
-            return_str += f"Batch {idx}: {len(data)} structure(s)\n"
-            for i, struct in enumerate(data):
-                return_str += f"Structure {i+1}: {struct.atom_array.array_length()} Atoms\n"
-                #return_str += struct.atom_array
-        return return_str
-    except Exception as e:
-        return f"Error: {str(e)}"
     
-
-# Initialize engine and run generation
-@spaces.GPU(duration=240) 
-def unconditional_generation(num_batches, num_designs_per_batch, length):
-    """
-    Runs an unconditional generation with the specified parameters for number of batches, number of designs per batch, and length of the generated proteins. 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.
-
-    Parameters:
-    ----------
-    num_batches: int or gr.Number,
-        The number of batches to generate.
-    num_designs_per_batch: int or gr.Number,
-        The number of designs to generate per batch.
-    length: int or gr.Number,
-        The length of the generated proteins.
-
-    Returns:
-    -------
-        directory: str,
-            The path to the directory where the generated structures are saved.
-        results: list of dicts,
-            A list of the generated structures' file paths, where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path".
-    """
-
-
-    config = RFD3InferenceConfig(
-    specification={
-        'length': length,
-    },
-    diffusion_batch_size=num_designs_per_batch,  # Generate 2 structures per batch
-    )
-
-    session_hash = gr.Request().session_hash
-    time_stamp = time.strftime("%Y-%m-%d-%H-%M-%S")
-    directory = f"./outputs/unconditional_generation/session_{session_hash}_{time_stamp}"
-    os.makedirs(directory, exist_ok=False)
-
-    try: 
-        model = RFD3InferenceEngine(**config)
-        outputs = model.run(
-            inputs=None,      # None for unconditional generation
-            out_dir=directory,     # None to return in memory (no file output)
-            n_batches=num_batches,      # Generate 1 batch
-        )
-
-        results = []
-        for batch in range(num_batches):
-            for design in range(num_designs_per_batch):
-                file_name = os.path.join(directory, f"_{batch}_model_{design}.cif.gz")
-                results.append({"batch": batch, "design": design, "cif_path": file_name, "pdb_path": mcif_gz_to_pdb(file_name)})
-            
-        print(results)
-        return directory, results
-    
-    except Exception as e:
-        raise RuntimeError(f"Error during generation: {str(e)}")
-    
-@spaces.GPU(duration=6120) 
-def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_per_batch, extra_args):
+#@spaces.GPU(duration=6120) 
+def generation_with_input_config_impl(input_file, pdb_file, num_batches, num_designs_per_batch, extra_args):
     """
     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.
 
@@ -166,4 +79,14 @@ def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_
     
     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 gr.update(value=f"Generation failed:\n{e.stderr}"), None, None
+
+
+def generation_with_input_config_factory(max_duration):
+
+    @spaces.GPU(duration=max_duration)
+    def generation_with_correct_time_limit(input_file, pdb_file, num_batches, num_designs_per_batch, extra_args):
+        return generation_with_input_config_impl(input_file, pdb_file, num_batches, num_designs_per_batch, extra_args)
+
+    return generation_with_correct_time_limit
+