Merge pull request #245 from llmsresearch/feat/multi-candidate

feat: multi-candidate generation with --num-candidates

Author: dippatel1994

README.md

@@ -263,6 +263,7 @@ paperbanana generate \
 | `--caption` | `-c` | Figure caption / communicative intent (required for new runs) |
 | `--output` | `-o` | Output image path (default: auto-generated in `outputs/`) |
 | `--iterations` | `-n` | Number of Visualizer-Critic refinement rounds (default: 3) |
+| `--num-candidates` | `-k` | Generate N candidate images in parallel, 1-8 (default: 1). Planning runs once; refinement fans out per candidate with seed offsets. Outputs land in `candidates/cand_<i>/`; the run-root `final_output` is candidate 1. Cost estimates and `--budget` account for the fan-out |
 | `--auto` | | Loop until critic is satisfied (with `--max-iterations` safety cap) |
 | `--max-iterations` | | Safety cap for `--auto` mode (default: 30) |
 | `--optimize` | | Preprocess inputs with parallel context enrichment and caption sharpening |

paperbanana/cli.py

@@ -314,6 +314,17 @@ def generate(
         "--budget",
         help="Budget cap in USD; pipeline aborts gracefully when exceeded",
     ),
+    num_candidates: Optional[int] = typer.Option(
+        None,
+        "--num-candidates",
+        "-k",
+        min=1,
+        max=8,
+        help=(
+            "Generate N candidate images in parallel (1-8). Planning runs once; "
+            "the Visualizer-Critic refinement fans out per candidate"
+        ),
+    ),
     dry_run: bool = typer.Option(
         False,
         "--dry-run",
@@ -514,6 +525,8 @@ def generate(
         overrides["seed"] = seed
     if budget is not None:
         overrides["budget_usd"] = budget
+    if num_candidates is not None:
+        overrides["num_candidates"] = num_candidates
     if venue:
         overrides["venue"] = venue
     if vector_export is not None:
@@ -728,6 +741,10 @@ async def _run_continue():
             f"Image: {settings.image_provider} / {settings.effective_image_model}",
             f"Iterations: {iter_est}",
             f"Optimize: {'yes' if settings.optimize_inputs else 'no'}",
+        ]
+        if settings.num_candidates > 1:
+            lines.append(f"Candidates: {settings.num_candidates} (parallel)")
+        lines += [
             "",
             f"Estimated VLM calls: {estimate['vlm_calls']}",
             f"Estimated image calls: {estimate['image_calls']}",
@@ -765,13 +782,16 @@ async def _run_continue():
     else:
         iter_label = str(settings.refinement_iterations)
 
+    candidates_label = (
+        f"\nCandidates: {settings.num_candidates} (parallel)" if settings.num_candidates > 1 else ""
+    )
     if not progress_json:
         console.print(
             Panel.fit(
                 f"[bold]PaperBanana[/bold] - Generating Methodology Diagram\n\n"
                 f"VLM: {settings.vlm_provider} / {settings.effective_vlm_model}\n"
                 f"Image: {settings.image_provider} / {settings.effective_image_model}\n"
-                f"Iterations: {iter_label}",
+                f"Iterations: {iter_label}{candidates_label}",
                 border_style="blue",
             )
         )
@@ -856,16 +876,19 @@ def on_progress(event: PipelineProgressEvent) -> None:
                         else " [green]✓[/green]"
                     )
             elif event.stage == PipelineProgressStage.VISUALIZER_START:
-                if event.iteration == 1:
-                    console.print("[bold]Phase 2[/bold] — Iterative Refinement")
                 extra = event.extra or {}
+                candidate = extra.get("candidate")
+                if event.iteration == 1 and candidate in (None, 1):
+                    console.print("[bold]Phase 2[/bold] — Iterative Refinement")
                 total = extra.get("total_iterations", 0)
                 if event.iteration and total:
                     label = f"{event.iteration}/{total}"
                 else:
                     label = str(event.iteration or "")
                 if settings.auto_refine:
                     label += " (auto)"
+                if candidate is not None:
+                    label = f"cand {candidate} · {label}"
                 console.print(f"  [dim]●[/dim] Generating image [{label}]...", end="")
             elif event.stage == PipelineProgressStage.VISUALIZER_END:
                 console.print(
@@ -927,6 +950,14 @@ def on_progress(event: PipelineProgressEvent) -> None:
         f" · {len(result.iterations)} iterations[/dim]\n"
     )
     console.print(f"  Output: [bold]{result.image_path}[/bold]")
+    for cand in result.metadata.get("candidates") or []:
+        if cand.get("error"):
+            console.print(
+                f"  Candidate {cand['index']}: [yellow]failed[/yellow] "
+                f"[dim]{str(cand['error'])[:120]}[/dim]"
+            )
+        elif cand.get("image_path"):
+            console.print(f"  Candidate {cand['index']}: {cand['image_path']}")
     if result.tikz_path:
         console.print(f"  TikZ:   [bold]{result.tikz_path}[/bold]")
     elif settings.export_tikz:

paperbanana/core/config.py

@@ -82,6 +82,12 @@ class Settings(BaseSettings):
     exemplar_retrieval_max_retries: int = 2
     venue: Venue = "neurips"
     vector_export: VectorExportMode = "none"
+    num_candidates: int = Field(
+        default=1,
+        ge=1,
+        le=8,
+        description="Number of parallel Phase-2 candidate branches (1-8)",
+    )
 
     # Reference settings
     reference_set_path: str = "data/reference_sets"
@@ -309,6 +315,7 @@ def _flatten_yaml(config: dict, prefix: str = "") -> dict:
         "pipeline.optimize_inputs": "optimize_inputs",
         "pipeline.output_resolution": "output_resolution",
         "pipeline.seed": "seed",
+        "pipeline.num_candidates": "num_candidates",
         "pipeline.exemplar_retrieval_enabled": "exemplar_retrieval_enabled",
         "pipeline.exemplar_retrieval_endpoint": "exemplar_retrieval_endpoint",
         "pipeline.exemplar_retrieval_mode": "exemplar_retrieval_mode",

paperbanana/core/cost_estimator.py

@@ -41,6 +41,8 @@ def estimate_cost(
     else:
         iterations = settings.refinement_iterations
 
+    num_candidates = max(1, getattr(settings, "num_candidates", 1))
+
     # Count expected API calls
     vlm_calls = 0
     image_calls = 0
@@ -75,10 +77,12 @@ def _image_cost() -> float:
     if diagram_type == DiagramType.METHODOLOGY and ve != "none":
         breakdown["structurer"] = _vlm_cost("structurer")
 
-    # Phase 2: Iterative refinement
+    # Phase 2: Iterative refinement. Multi-candidate fan-out runs Phase 2
+    # once per candidate (Phase 1 planning is shared), so visualizer and
+    # critic costs scale by num_candidates.
     vis_total = 0.0
     critic_total = 0.0
-    for _ in range(iterations):
+    for _ in range(iterations * num_candidates):
         if diagram_type == DiagramType.STATISTICAL_PLOT:
             vis_total += _vlm_cost("visualizer_vlm")
         else:
@@ -98,11 +102,17 @@ def _image_cost() -> float:
             f"Auto-refine: estimated for max {iterations} iterations; "
             "actual cost may be lower if critic is satisfied early"
         )
+    if num_candidates > 1:
+        notes.append(
+            f"Multi-candidate: visualizer/critic costs scaled by "
+            f"{num_candidates} parallel candidates"
+        )
 
     return {
         "estimated_total_usd": round(total, 6),
         "vlm_calls": vlm_calls,
         "image_calls": image_calls,
+        "num_candidates": num_candidates,
         "breakdown_by_agent": {k: round(v, 6) for k, v in breakdown.items()},
         "pricing_note": "; ".join(notes) if notes else None,
     }