update

app.py

@@ -15,6 +15,7 @@ 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)
@@ -38,6 +39,7 @@ 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')
 
 
@@ -172,10 +174,7 @@ def generate(input_file):
     return [
         output.IFRAME_TEMPLATE.format(html=html),
         [inp_sdf] + out_files,
-        gr.Radio.update(
-            choices=['Sample 1', 'Sample 2', 'Sample 3', 'Sample 4', 'Sample 5'],
-            value='Sample 1',
-        )
+        gr.Radio.update(visible=True)
     ]
 
 
@@ -203,7 +202,14 @@ with demo:
                 gr.Markdown('## Visualization')
                 gr.Markdown('Below you will see input and output molecules')
                 visualization = gr.HTML()
-                samples = gr.Radio(interactive=True, type='index', label='Samples')
+                samples = gr.Radio(
+                    choices=['Sample 1', 'Sample 2', 'Sample 3', 'Sample 4', 'Sample 5'],
+                    value='Sample 1',
+                    type='index',
+                    show_label=False,
+                    visible=False,
+                    interactive=True,
+                )
 
     input_file.change(
         fn=show_input,