Change max batch_size

app.py

@@ -17,6 +17,11 @@ from src.generation import generate_linkers, try_to_convert_to_sdf, get_pocket
 from zipfile import ZipFile
 
 
+MIN_N_STEPS = 100
+MAX_N_STEPS = 500
+MAX_BATCH_SIZE = 5
+
+
 MODELS_METADATA = {
     'geom_difflinker': {
         'link': 'https://zenodo.org/record/7121300/files/geom_difflinker.ckpt?download=1',
@@ -329,9 +334,7 @@ def generate_without_pocket(input_file, n_steps, n_atoms, num_samples, selected_
 
     for data in dataloader:
         try:
-            generate_linkers(
-                ddpm=ddpm, data=data, num_samples=num_samples, sample_fn=sample_fn, name=name, with_pocket=False
-            )
+            generate_linkers(ddpm=ddpm, data=data, sample_fn=sample_fn, name=name, with_pocket=False)
         except Exception as e:
             e = str(e).replace('\'', '')
             error = f'Caught exception while generating linkers: {e}'
@@ -450,7 +453,8 @@ def generate_with_pocket(in_fragments, in_protein, n_steps, n_atoms, num_samples
     dataset = MOADDataset(data=dataset)
     ddpm.val_dataset = dataset
 
-    dataloader = get_dataloader(dataset, batch_size=num_samples, collate_fn=collate_with_fragment_edges)
+    batch_size = min(num_samples, MAX_BATCH_SIZE)
+    dataloader = get_dataloader(dataset, batch_size=batch_size, collate_fn=collate_with_fragment_edges)
     print('Created dataloader')
 
     ddpm.edm.T = n_steps
@@ -470,10 +474,13 @@ def generate_with_pocket(in_fragments, in_protein, n_steps, n_atoms, num_samples
         def sample_fn(_data):
             return torch.ones(_data['positions'].shape[0], device=device, dtype=torch.long) * n_atoms
 
-    for data in dataloader:
+    for batch_i, data in enumerate(dataloader):
         try:
+            offset_idx = batch_i * batch_size
             generate_linkers(
-                ddpm=ddpm, data=data, num_samples=num_samples, sample_fn=sample_fn, name=name, with_pocket=True
+                ddpm=ddpm, data=data,
+                sample_fn=sample_fn, name=name, with_pocket=True,
+                offset_idx=offset_idx,
             )
         except Exception as e:
             e = str(e).replace('\'', '')
@@ -520,7 +527,10 @@ with demo:
                 gr.Markdown('Upload the file of the target protein in .pdb format (optionally):')
                 input_protein_file = gr.File(file_count='single', label='Target Protein (Optional)')
 
-                n_steps = gr.Slider(minimum=50, maximum=500, label="Number of Denoising Steps", step=10)
+                n_steps = gr.Slider(
+                    minimum=MIN_N_STEPS, maximum=MAX_N_STEPS,
+                    label="Number of Denoising Steps", step=10
+                )
                 n_atoms = gr.Slider(
                     minimum=0, maximum=20,
                     label="Linker Size: DiffLinker will predict it if set to 0",

src/generation.py

@@ -10,7 +10,7 @@ from src.utils import FoundNaNException
 from src.datasets import get_one_hot
 
 
-def generate_linkers(ddpm, data, num_samples, sample_fn, name, with_pocket=False):
+def generate_linkers(ddpm, data, sample_fn, name, with_pocket=False, offset_idx=0):
     chain = node_mask = None
     for i in range(5):
         try:
@@ -37,7 +37,8 @@ def generate_linkers(ddpm, data, num_samples, sample_fn, name, with_pocket=False
     if with_pocket:
         node_mask[torch.where(data['pocket_mask'])] = 0
 
-    names = [f'output_{i + 1}_{name}' for i in range(num_samples)]
+    batch_size = len(data)
+    names = [f'output_{offset_idx + i + 1}_{name}' for i in range(batch_size)]
     save_xyz_file('results', h, x, node_mask, names=names, is_geom=True, suffix='')
     print('Saved XYZ files')