now supports repeated generation runs with fixed batch size per run:

app.py

@@ -315,6 +315,7 @@ def save_to_format(xyz_str, idx, fmt="pdb"):
 def generate(
     model_name,
     num_molecules,
+    num_runs,
     size_mode,
     fixed_size,
     diffusion_steps,
@@ -333,7 +334,11 @@ def generate(
     if TASK is None:
         return "", gr.update(choices=[], value=None), [], None, None
     
-    print(f"Generating {num_molecules} molecules with '{model_name}' (Steps: {diffusion_steps}, Seed: {seed})...")
+    total_requested = int(num_molecules) * int(num_runs)
+    print(
+        f"Generating {total_requested} molecules as {num_runs} run(s) x batch {num_molecules} "
+        f"with '{model_name}' (Steps: {diffusion_steps}, Seed: {seed})..."
+    )
     seed_everything(int(seed))
     
     # 1. Override diffusion steps
@@ -347,21 +352,7 @@ def generate(
         TASK.T = int(diffusion_steps)
         
     try:
-        # 2. Determine molecule sizes
-        # Handle "Auto" string check robustly
-        if "Auto" in size_mode:
-            if TASK.node_dist_model is not None:
-                # DistributionNodes.sample returns tensor of sizes
-                nodesxsample = TASK.node_dist_model.sample(num_molecules)
-            else:
-                # Fallback if no dist model
-                nodesxsample = torch.randint(10, 30, (num_molecules,))
-        else:
-            nodesxsample = torch.tensor([fixed_size] * num_molecules)
-            
-        nodesxsample = nodesxsample.to(DEVICE).long()
-        
-        # 3. Sample (auto-switch unconditional/conditional)
+        # 2. Resolve conditioning mode/inputs once
         condition_names = get_condition_names(TASK)
         is_conditional = len(condition_names) > 0
 
@@ -372,45 +363,58 @@ def generate(
             negative_values = parse_condition_row(
                 negative_values_df, condition_names, required=False
             )
-            target_fn = lambda z, t: torch.zeros(
-                z.size(0), device=z.device, dtype=z.dtype
-            )
-            one_hot, charges, x, node_mask = TASK.sample_guidance_conitional(
-                target_function=target_fn,
-                target_value=target_values,
-                negative_target_value=negative_values,
-                nodesxsample=nodesxsample,
-                gg_scale=0.0,
-                cfg_scale=float(cfg_scale),
-                guidance_ver="cfg",
-                n_frames=0,
-                fix_noise=False,
-            )
-        else:
-            # This returns: one_hot [B,N,C], charges [B,N,1], x [B,N,3], node_mask [B,N,1]
-            one_hot, charges, x, node_mask = TASK.sample(
-                nodesxsample=nodesxsample,
-                mode="ddpm",
-                n_frames=0,
-                fix_noise=False,
-            )
-        
-        # 4. Post-process
+
+        # 3. Sample across runs
         xyz_strings = []
         summary_rows = []
-        
-        for i in range(num_molecules):
-            # Extract single molecule data
-            xyz_str = tensors_to_xyz_string(
-                one_hot[i], 
-                x[i], 
-                node_mask[i], 
-                TASK.atom_vocab
-            )
-            xyz_strings.append(xyz_str)
-            summary_rows.append(parse_composition(xyz_str))
+
+        target_fn = lambda z, t: torch.zeros(
+            z.size(0), device=z.device, dtype=z.dtype
+        )
+
+        for run_idx in range(int(num_runs)):
+            # Determine molecule sizes for this run
+            if "Auto" in size_mode:
+                if TASK.node_dist_model is not None:
+                    nodesxsample = TASK.node_dist_model.sample(int(num_molecules))
+                else:
+                    nodesxsample = torch.randint(10, 30, (int(num_molecules),))
+            else:
+                nodesxsample = torch.tensor([fixed_size] * int(num_molecules))
+
+            nodesxsample = nodesxsample.to(DEVICE).long()
+
+            if is_conditional:
+                one_hot, charges, x, node_mask = TASK.sample_guidance_conitional(
+                    target_function=target_fn,
+                    target_value=target_values,
+                    negative_target_value=negative_values,
+                    nodesxsample=nodesxsample,
+                    gg_scale=0.0,
+                    cfg_scale=float(cfg_scale),
+                    guidance_ver="cfg",
+                    n_frames=0,
+                    fix_noise=False,
+                )
+            else:
+                one_hot, charges, x, node_mask = TASK.sample(
+                    nodesxsample=nodesxsample,
+                    mode="ddpm",
+                    n_frames=0,
+                    fix_noise=False,
+                )
+
+            for i in range(int(num_molecules)):
+                xyz_str = tensors_to_xyz_string(
+                    one_hot[i],
+                    x[i],
+                    node_mask[i],
+                    TASK.atom_vocab
+                )
+                xyz_strings.append(xyz_str)
+                summary_rows.append(parse_composition(xyz_str))
             
-        # 5. Output generation – save zip for bulk download
+        # 4. Output generation – save zip for bulk download
         zip_path = create_xyz_zip(xyz_strings)
         
         # Prepare table with "Name" column
@@ -746,6 +750,7 @@ with gr.Blocks(title="MolCraftDiffusion", theme=gr.themes.Soft(), head=THREEDMOL
                     "Mode: **Conditional**" if initial_is_conditional else "Mode: **Unconditional**"
                 )
                 num_mol = gr.Slider(1, 12, value=4, step=1, label="Number of Molecules")
+                num_runs = gr.Slider(1, 20, value=1, step=1, label="Number of Runs")
                 
                 size_mode = gr.Radio(
                     ["Auto (from training data)", "Fixed size"],
@@ -1035,6 +1040,7 @@ with gr.Blocks(title="MolCraftDiffusion", theme=gr.themes.Soft(), head=THREEDMOL
         inputs=[
             model_selector,
             num_mol,
+            num_runs,
             size_mode,
             fixed_size,
             diffusion_steps,