n_steps

app.py

@@ -15,7 +15,6 @@ from src.lightning import DDPM
 from src.linker_size_lightning import SizeClassifier
 
 N_SAMPLES = 5
-N_STEPS = 10
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--ip', type=str, default=None)
@@ -39,7 +38,6 @@ if not os.path.exists(diffusion_path):
     link = 'https://zenodo.org/record/7121300/files/geom_difflinker.ckpt?download=1'
     subprocess.run(f'wget {link} -O {diffusion_path}', shell=True)
 ddpm = DDPM.load_from_checkpoint('models/geom_difflinker.ckpt', map_location=device).eval().to(device)
-ddpm.edm.T = N_STEPS
 print('Loaded diffusion model')
 
 
@@ -111,7 +109,7 @@ def draw_sample(idx, out_files):
     return output.IFRAME_TEMPLATE.format(html=html)
 
 
-def generate(input_file):
+def generate(input_file, n_steps):
     if input_file is None:
         return ''
 
@@ -155,6 +153,8 @@ def generate(input_file):
     dataloader = get_dataloader(dataset, batch_size=N_SAMPLES, collate_fn=collate_with_fragment_edges)
     print('Created dataloader')
 
+    ddpm.edm.T = n_steps
+
     for data in dataloader:
         chain, node_mask = ddpm.sample_chain(data, sample_fn=sample_fn, keep_frames=1)
         print('Generated linker')
@@ -188,7 +188,8 @@ with demo:
             with gr.Column():
                 gr.Markdown('## Input Fragments')
                 gr.Markdown('Upload the file with 3D-coordinates of the input fragments in .pdb, .mol2 or .sdf format:')
-                input_file = gr.File(file_count='single', label='Input Fragments') 
+                input_file = gr.File(file_count='single', label='Input Fragments')
+                n_steps = gr.Slider(minimum=10, maximum=500, label="Number of Diffusion Steps", step=10)
                 examples = gr.Dataset(
                     components=[gr.File(visible=False)],
                     samples=[['examples/example_1.sdf'], ['examples/example_2.sdf']],
@@ -219,13 +220,13 @@ with demo:
         outputs=[visualization],
     )
     examples.click(
-        fn=lambda idx: [f'examples/example_{idx+1}.sdf', show_input(f'examples/example_{idx+1}.sdf')],
+        fn=lambda idx: [f'examples/example_{idx+1}.sdf', 10, show_input(f'examples/example_{idx+1}.sdf')],
         inputs=[examples],
-        outputs=[input_file, visualization]
+        outputs=[input_file, n_steps, visualization]
     )
     button.click(
         fn=generate,
-        inputs=[input_file],
+        inputs=[input_file, n_steps],
         outputs=[visualization, output_files, samples],
     )
     samples.change(