NNGen/README.md

Multi-Agent Neural Network Diagram Generator (Skeleton) — Gemini 2.5 Flash Image

This repository is a minimal, runnable skeleton that turns a textual NN spec into a publication-style diagram via a multi-agent pipeline:

• Parser → Planner → Prompt-Generator → Image-Generator (G1) → Label-Generator (G2) → Judge → Selector → (Editor loop) → Archivist
• All model calls flow through call_gemini(...), making it easy to use Gemini 2.5 Flash for text and Gemini 2.5 Flash Image for images.

Key additions in this version

• Two-stage generation: G1 draws the geometry-only skeleton (no text), G2 overlays labels on top of the skeleton.
• Hard violations: Judge returns actionable violations; missing labels are flagged as HARD to trigger edits reliably.
• Parallelism: G1, G2, and Judge run in parallel; set NNG_CONCURRENCY (default 4).
• Remote images by default: image generate/edit use Gemini 2.5 Flash Image models. If API is missing, the system can fall back to a local placeholder to stay runnable.

Quick Start

1. Python 3.10+

2. Install deps

pip install -r requirements.txt

3. Configure Gemini (choose one)

• Env var: export GEMINI_API_KEY=YOUR_KEY
• File: create app/llm/credentials.py with GEMINI_API_KEY = "YOUR_KEY"

4. Run (K=candidates, T=max edit rounds)

# Text mode (spec -> image)
python -m app.cli --mode text --spec spec/vit.txt --K 4 --T 1

# Image mode (text + image fusion/edit)
# Example: edit an existing diagram with a component replacement using a reference image
python -m app.cli --mode image --base-image path/to/base.png \
  --ref-image path/to/transformer_ref.png \
  --instructions "Replace the UNet backbone with a Transformer (DiT); keep layout, font, and colors consistent."

Artifacts are saved under artifacts/run_YYYYmmdd_HHMMSS/ with final.png as the chosen result.

Gemini 2.5 Flash Image in This Project

• G1 geometry: gen_generate.py calls GEMINI_IMAGE_MODEL (Gemini 2.5 Flash Image) to render a clean, geometry-only skeleton quickly.
• G2 labels: gen_labels.py uses GEMINI_IMAGE_EDIT_MODEL to overlay text labels onto the G1 skeleton without redrawing everything.
• Edit loop: edit.py performs targeted corrections via the same image model, enabling fast, iterative refinements instead of full regenerations.
• Why it matters: the model’s speed and editability make multi-round diagram refinement practical while preserving layout quality.
• Fallback: if no API key is available, the pipeline remains runnable using local placeholders generated by app/llm/gemini.py.

Models

• GEMINI_MODEL (default gemini-2.5-flash): parsing, planning, prompt generation, and judging.
• GEMINI_IMAGE_MODEL (recommended gemini-2.5-flash-image or gemini-2.5-flash-image-preview): image generation (G1).
• GEMINI_IMAGE_EDIT_MODEL (recommended gemini-2.5-flash-image or gemini-2.5-flash-image-preview): image editing (G2, Editor). Notes: If GEMINI_API_KEY is not set, the pipeline uses offline placeholders to remain runnable. With an API key present, you must set valid image model env vars; errors are raised if image models are unset or calls fail (no automatic local fallback).

Fusion Mode (Text + Image)

• Accepts a base diagram (--base-image) and optional reference images (--ref-image repeatable) plus instructions.
• Uses Gemini 2.5 Flash Image to compose images under textual guidance – ideal for swapping a module (e.g., UNet → Transformer) while preserving style and layout.
• Outputs multiple fused candidates (K) and archives the first as final.png.

Structure

app/
  cli.py              # CLI entry (K/T/outdir)
  graph.py            # Orchestrator + edit loop
  state.py            # AppState + artifacts
  prompts.py          # Centralized prompts (parse/plan/G1/G2/judge/edit)
  nodes/
    parser.py, planner.py, prompt_gen.py
    gen_generate.py   # G1 skeleton images (no text)
    gen_labels.py     # G2 label overlay edits
    judge.py, select.py, edit.py, archive.py
  llm/
    gemini.py         # Unified wrapper (API + offline fallback)
    credentials.example.py
spec/
  vit.txt             # Example ViT spec (English)
artifacts/            # Outputs per run

Tips

• Concurrency: NNG_CONCURRENCY=4 python -m app.cli --spec ...
• Tuning: Start with K=4, T=1; increase T for more correction rounds.
• Debug: image calls write *.resp.txt/*.meta.json alongside outputs (can be removed later if undesired).