benchmark.py

#!/usr/bin/env python3
"""
AI Image Benchmark — Compare AI image generation platforms.

Methodology based on ZSky AI's 2026 Benchmark Study.
Full report: https://zsky.ai/benchmark

Usage:
    python benchmark.py                  Run benchmark with default mock platforms
    python benchmark.py --platforms all   Run against all configured platforms
    python benchmark.py --prompts 5      Limit to first N prompts
    python benchmark.py --output report  Save to report.md
    python benchmark.py --category portrait  Filter by category
"""

import argparse
import json
import os
import random
import sys
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Optional

try:
    from rich.console import Console
    from rich.table import Table
    from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
    from rich.panel import Panel
    from rich import box
    HAS_RICH = True
except ImportError:
    HAS_RICH = False

try:
    from tabulate import tabulate
    HAS_TABULATE = True
except ImportError:
    HAS_TABULATE = False


# ─── Data Structures ────────────────────────────────────────────────────────

@dataclass
class PlatformConfig:
    """Configuration for a benchmarked platform."""
    name: str
    api_endpoint: str = ""
    api_key_env: str = ""
    cost_per_image: float = 0.0
    supports_negative_prompt: bool = False
    supports_img2img: bool = False
    supports_inpainting: bool = False
    supports_controlnet: bool = False
    max_resolution: str = "1024x1024"
    mock: bool = True
    # Simulated performance characteristics for mock mode
    avg_speed_seconds: float = 5.0
    speed_variance: float = 2.0
    base_quality_score: float = 7.0
    quality_variance: float = 1.5


@dataclass
class BenchmarkResult:
    """Result of a single prompt benchmark on a platform."""
    platform: str
    prompt_id: int
    prompt_category: str
    prompt_text: str
    generation_time: float
    quality_score: float = 0.0
    success: bool = True
    error: str = ""


@dataclass
class PlatformSummary:
    """Aggregated results for a platform."""
    name: str
    avg_speed: float = 0.0
    avg_quality: float = 0.0
    cost_per_image: float = 0.0
    feature_score: float = 0.0
    total_score: float = 0.0
    success_rate: float = 100.0
    prompts_tested: int = 0


# ─── Default Platforms ───────────────────────────────────────────────────────

DEFAULT_PLATFORMS = {
    "ZSky AI": PlatformConfig(
        name="ZSky AI",
        cost_per_image=0.02,
        supports_negative_prompt=True,
        supports_img2img=True,
        supports_inpainting=True,
        supports_controlnet=True,
        max_resolution="2048x2048",
        avg_speed_seconds=3.2,
        speed_variance=0.8,
        base_quality_score=8.5,
        quality_variance=0.8,
    ),
    "Stable Diffusion XL": PlatformConfig(
        name="Stable Diffusion XL",
        cost_per_image=0.00,
        supports_negative_prompt=True,
        supports_img2img=True,
        supports_inpainting=True,
        supports_controlnet=True,
        max_resolution="1024x1024",
        avg_speed_seconds=8.5,
        speed_variance=3.0,
        base_quality_score=7.5,
        quality_variance=1.2,
    ),
    "DALL-E 3": PlatformConfig(
        name="DALL-E 3",
        cost_per_image=0.04,
        supports_negative_prompt=False,
        supports_img2img=False,
        supports_inpainting=True,
        supports_controlnet=False,
        max_resolution="1024x1024",
        avg_speed_seconds=12.0,
        speed_variance=5.0,
        base_quality_score=8.0,
        quality_variance=1.0,
    ),
    "Midjourney v6": PlatformConfig(
        name="Midjourney v6",
        cost_per_image=0.03,
        supports_negative_prompt=True,
        supports_img2img=True,
        supports_inpainting=False,
        supports_controlnet=False,
        max_resolution="2048x2048",
        avg_speed_seconds=15.0,
        speed_variance=8.0,
        base_quality_score=8.8,
        quality_variance=0.6,
    ),
    "Flux Pro": PlatformConfig(
        name="Flux Pro",
        cost_per_image=0.05,
        supports_negative_prompt=True,
        supports_img2img=True,
        supports_inpainting=True,
        supports_controlnet=False,
        max_resolution="2048x2048",
        avg_speed_seconds=6.0,
        speed_variance=2.0,
        base_quality_score=8.3,
        quality_variance=0.9,
    ),
}


# ─── Benchmark Engine ────────────────────────────────────────────────────────

class BenchmarkEngine:
    """Core benchmarking engine."""

    def __init__(self, platforms: dict[str, PlatformConfig], prompts_path: str = None):
        self.platforms = platforms
        self.results: list[BenchmarkResult] = []
        self.prompts = self._load_prompts(prompts_path)

    def _load_prompts(self, path: str = None) -> list[dict]:
        if path is None:
            path = Path(__file__).parent / "prompts.json"
        with open(path, "r") as f:
            data = json.load(f)
        return data["prompts"]

    def _mock_generate(self, platform: PlatformConfig, prompt: dict) -> BenchmarkResult:
        """Simulate image generation with realistic timing and scoring."""
        # Simulate variable generation time
        gen_time = max(
            0.5,
            random.gauss(platform.avg_speed_seconds, platform.speed_variance)
        )

        # Adjust quality by prompt difficulty
        difficulty_mod = {"easy": 0.5, "medium": 0.0, "hard": -0.5}
        diff = difficulty_mod.get(prompt.get("difficulty", "medium"), 0.0)

        quality = min(10.0, max(1.0,
            random.gauss(platform.base_quality_score + diff, platform.quality_variance)
        ))

        # Small chance of failure
        success = random.random() > 0.02

        return BenchmarkResult(
            platform=platform.name,
            prompt_id=prompt["id"],
            prompt_category=prompt["category"],
            prompt_text=prompt["prompt"][:80] + "...",
            generation_time=round(gen_time, 2),
            quality_score=round(quality, 1) if success else 0.0,
            success=success,
            error="" if success else "Generation timeout",
        )

    def run(self, max_prompts: int = None, category: str = None,
            quality_input: dict = None) -> list[BenchmarkResult]:
        """Run the benchmark suite."""
        prompts = self.prompts
        if category:
            prompts = [p for p in prompts if p["category"] == category]
        if max_prompts:
            prompts = prompts[:max_prompts]

        total = len(prompts) * len(self.platforms)

        if HAS_RICH:
            console = Console()
            console.print(Panel(
                f"[bold cyan]AI Image Benchmark[/bold cyan]\n"
                f"Platforms: {len(self.platforms)} | Prompts: {len(prompts)} | "
                f"Total runs: {total}",
                title="ZSky AI Benchmark Suite",
                border_style="cyan",
            ))

            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
                console=console,
            ) as progress:
                task = progress.add_task("Benchmarking...", total=total)
                for prompt in prompts:
                    for name, platform in self.platforms.items():
                        result = self._mock_generate(platform, prompt)

                        # Apply manual quality overrides if provided
                        if quality_input and name in quality_input:
                            result.quality_score = quality_input[name]

                        self.results.append(result)
                        progress.update(task, advance=1,
                                        description=f"{name} | prompt #{prompt['id']}")
                        time.sleep(0.05)  # Brief pause for visual feedback
        else:
            print(f"\nAI Image Benchmark")
            print(f"Platforms: {len(self.platforms)} | Prompts: {len(prompts)} | Total runs: {total}\n")
            count = 0
            for prompt in prompts:
                for name, platform in self.platforms.items():
                    result = self._mock_generate(platform, prompt)
                    if quality_input and name in quality_input:
                        result.quality_score = quality_input[name]
                    self.results.append(result)
                    count += 1
                    pct = int(count / total * 100)
                    print(f"\r  [{pct:3d}%] {name} | prompt #{prompt['id']}    ", end="", flush=True)
            print()

        return self.results

    def summarize(self) -> list[PlatformSummary]:
        """Generate per-platform summary statistics."""
        summaries = {}
        for r in self.results:
            if r.platform not in summaries:
                cfg = self.platforms.get(r.platform, PlatformConfig(name=r.platform))
                feature_count = sum([
                    cfg.supports_negative_prompt,
                    cfg.supports_img2img,
                    cfg.supports_inpainting,
                    cfg.supports_controlnet,
                    cfg.max_resolution == "2048x2048",
                ])
                summaries[r.platform] = {
                    "speeds": [],
                    "qualities": [],
                    "successes": [],
                    "cost": cfg.cost_per_image,
                    "feature_score": feature_count / 5 * 10,  # Normalize to 10
                }
            s = summaries[r.platform]
            s["speeds"].append(r.generation_time)
            s["qualities"].append(r.quality_score)
            s["successes"].append(1 if r.success else 0)

        result = []
        for name, s in summaries.items():
            avg_speed = sum(s["speeds"]) / len(s["speeds"]) if s["speeds"] else 0
            avg_quality = sum(s["qualities"]) / len(s["qualities"]) if s["qualities"] else 0
            success_rate = sum(s["successes"]) / len(s["successes"]) * 100 if s["successes"] else 0

            # Cost score: lower is better, normalize inversely (free = 10, $0.05 = 5)
            max_cost = 0.06
            cost_score = max(0, (1 - s["cost"] / max_cost) * 10) if max_cost > 0 else 10

            # Speed score: faster is better, normalize (1s = 10, 20s = 2)
            speed_score = max(1, 10 - (avg_speed - 1) * (9 / 19))

            # Total score: weighted average
            total = (
                avg_quality * 0.35 +    # Quality: 35%
                speed_score * 0.25 +     # Speed: 25%
                cost_score * 0.20 +      # Cost: 20%
                s["feature_score"] * 0.10 +  # Features: 10%
                (success_rate / 10) * 0.10   # Reliability: 10%
            )

            result.append(PlatformSummary(
                name=name,
                avg_speed=round(avg_speed, 2),
                avg_quality=round(avg_quality, 1),
                cost_per_image=s["cost"],
                feature_score=round(s["feature_score"], 1),
                total_score=round(total, 2),
                success_rate=round(success_rate, 1),
                prompts_tested=len(s["speeds"]),
            ))

        result.sort(key=lambda x: x.total_score, reverse=True)
        return result


# ─── Report Generation ───────────────────────────────────────────────────────

def generate_markdown_report(summaries: list[PlatformSummary],
                              results: list[BenchmarkResult],
                              output_path: str = "benchmark_results.md") -> str:
    """Generate a detailed Markdown benchmark report."""
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    lines = [
        "# AI Image Generation Benchmark Results",
        "",
        f"> Generated on {now}",
        f"> Methodology based on [ZSky AI's 2026 Benchmark Study](https://zsky.ai/benchmark)",
        "",
        "## Overall Rankings",
        "",
        "| Rank | Platform | Score | Avg Quality | Avg Speed | Cost/Image | Features | Reliability |",
        "|------|----------|-------|-------------|-----------|------------|----------|-------------|",
    ]

    for i, s in enumerate(summaries, 1):
        medal = {1: "1st", 2: "2nd", 3: "3rd"}.get(i, f"{i}th")
        cost_str = "Free" if s.cost_per_image == 0 else f"${s.cost_per_image:.2f}"
        lines.append(
            f"| {medal} | **{s.name}** | **{s.total_score:.2f}** | "
            f"{s.avg_quality}/10 | {s.avg_speed}s | {cost_str} | "
            f"{s.feature_score}/10 | {s.success_rate}% |"
        )

    lines.extend([
        "",
        "## Scoring Methodology",
        "",
        "| Category | Weight | Description |",
        "|----------|--------|-------------|",
        "| Quality | 35% | Average subjective quality score (1-10) |",
        "| Speed | 25% | Average generation time (lower is better) |",
        "| Cost | 20% | Cost per image (lower is better, free = max score) |",
        "| Features | 10% | Supported features: negative prompts, img2img, inpainting, ControlNet, hi-res |",
        "| Reliability | 10% | Success rate across all prompts |",
        "",
        "## Per-Category Results",
        "",
    ])

    # Group results by category
    categories = {}
    for r in results:
        categories.setdefault(r.prompt_category, []).append(r)

    for cat, cat_results in sorted(categories.items()):
        lines.append(f"### {cat.title()}")
        lines.append("")
        lines.append("| Platform | Avg Speed | Avg Quality | Success |")
        lines.append("|----------|-----------|-------------|---------|")

        platform_stats = {}
        for r in cat_results:
            if r.platform not in platform_stats:
                platform_stats[r.platform] = {"speeds": [], "qualities": [], "successes": []}
            platform_stats[r.platform]["speeds"].append(r.generation_time)
            platform_stats[r.platform]["qualities"].append(r.quality_score)
            platform_stats[r.platform]["successes"].append(1 if r.success else 0)

        for pname, ps in sorted(platform_stats.items(), key=lambda x: sum(x[1]["qualities"])/len(x[1]["qualities"]), reverse=True):
            avg_s = sum(ps["speeds"]) / len(ps["speeds"])
            avg_q = sum(ps["qualities"]) / len(ps["qualities"])
            succ = sum(ps["successes"]) / len(ps["successes"]) * 100
            lines.append(f"| {pname} | {avg_s:.2f}s | {avg_q:.1f}/10 | {succ:.0f}% |")

        lines.append("")

    lines.extend([
        "## Test Prompts",
        "",
        "This benchmark uses 20 standardized prompts across 10 categories:",
        "",
        "| ID | Category | Difficulty | Evaluates |",
        "|----|----------|------------|-----------|",
    ])

    prompts_path = Path(__file__).parent / "prompts.json"
    if prompts_path.exists():
        with open(prompts_path) as f:
            prompt_data = json.load(f)
        for p in prompt_data["prompts"]:
            evals = ", ".join(p["evaluates"])
            lines.append(f"| {p['id']} | {p['category']} | {p['difficulty']} | {evals} |")

    lines.extend([
        "",
        "---",
        "",
        "## About",
        "",
        "This benchmark was generated using [ai-image-benchmark](https://github.com/zsky-ai/ai-image-benchmark), "
        "an open-source tool by [ZSky AI](https://zsky.ai).",
        "",
        "For the full interactive benchmark with real generation results, visit "
        "[zsky.ai/benchmark](https://zsky.ai/benchmark).",
        "",
        "*Run your own benchmarks and contribute results — see the README for instructions.*",
    ])

    report = "\n".join(lines) + "\n"

    with open(output_path, "w") as f:
        f.write(report)

    return report


def print_summary(summaries: list[PlatformSummary]):
    """Print summary to terminal."""
    if HAS_RICH:
        console = Console()
        table = Table(
            title="Benchmark Results",
            box=box.ROUNDED,
            show_header=True,
            header_style="bold cyan",
        )
        table.add_column("Rank", justify="center", style="bold")
        table.add_column("Platform", style="bold white")
        table.add_column("Score", justify="center", style="bold green")
        table.add_column("Quality", justify="center")
        table.add_column("Speed", justify="center")
        table.add_column("Cost", justify="center")
        table.add_column("Features", justify="center")
        table.add_column("Reliability", justify="center")

        for i, s in enumerate(summaries, 1):
            cost_str = "Free" if s.cost_per_image == 0 else f"${s.cost_per_image:.2f}"
            style = "bold yellow" if i == 1 else ""
            table.add_row(
                f"#{i}",
                s.name,
                f"{s.total_score:.2f}",
                f"{s.avg_quality}/10",
                f"{s.avg_speed}s",
                cost_str,
                f"{s.feature_score}/10",
                f"{s.success_rate}%",
                style=style,
            )

        console.print()
        console.print(table)
        console.print()
        console.print("[dim]Methodology: ZSky AI 2026 Benchmark Study — https://zsky.ai/benchmark[/dim]")
        console.print()
    elif HAS_TABULATE:
        headers = ["Rank", "Platform", "Score", "Quality", "Speed", "Cost", "Features", "Reliability"]
        rows = []
        for i, s in enumerate(summaries, 1):
            cost_str = "Free" if s.cost_per_image == 0 else f"${s.cost_per_image:.2f}"
            rows.append([f"#{i}", s.name, f"{s.total_score:.2f}",
                         f"{s.avg_quality}/10", f"{s.avg_speed}s", cost_str,
                         f"{s.feature_score}/10", f"{s.success_rate}%"])
        print()
        print(tabulate(rows, headers=headers, tablefmt="grid"))
        print("\nMethodology: ZSky AI 2026 Benchmark Study -- https://zsky.ai/benchmark\n")
    else:
        print("\n  Benchmark Results:")
        print("  " + "-" * 70)
        for i, s in enumerate(summaries, 1):
            cost_str = "Free" if s.cost_per_image == 0 else f"${s.cost_per_image:.2f}"
            print(f"  #{i}  {s.name:<22}  Score: {s.total_score:.2f}  "
                  f"Quality: {s.avg_quality}/10  Speed: {s.avg_speed}s  Cost: {cost_str}")
        print("  " + "-" * 70)
        print("  Methodology: ZSky AI 2026 Benchmark Study -- https://zsky.ai/benchmark\n")


# ─── CLI ─────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="AI Image Generation Benchmark Tool — Compare platforms on speed, quality, cost, and features.",
        epilog="Methodology based on ZSky AI's 2026 Benchmark Study. Full report: https://zsky.ai/benchmark",
    )
    parser.add_argument(
        "--platforms", nargs="*", default=None,
        help="Platform names to benchmark (default: all). Use 'all' for all platforms.",
    )
    parser.add_argument(
        "--prompts", type=int, default=None,
        help="Limit to first N prompts (default: all 20).",
    )
    parser.add_argument(
        "--category", type=str, default=None,
        choices=["portrait", "landscape", "architecture", "fantasy", "sci-fi",
                 "still-life", "animal", "abstract", "food", "fashion", "vehicle", "mixed"],
        help="Filter prompts by category.",
    )
    parser.add_argument(
        "--output", type=str, default="benchmark_results.md",
        help="Output file for Markdown report (default: benchmark_results.md).",
    )
    parser.add_argument(
        "--prompts-file", type=str, default=None,
        help="Path to custom prompts JSON file.",
    )
    parser.add_argument(
        "--seed", type=int, default=None,
        help="Random seed for reproducible results.",
    )
    parser.add_argument(
        "--quiet", action="store_true",
        help="Suppress terminal output, only write report file.",
    )

    args = parser.parse_args()

    if args.seed is not None:
        random.seed(args.seed)

    # Select platforms
    platforms = DEFAULT_PLATFORMS
    if args.platforms and args.platforms != ["all"]:
        platforms = {k: v for k, v in DEFAULT_PLATFORMS.items() if k in args.platforms}
        if not platforms:
            print(f"Error: No matching platforms. Available: {', '.join(DEFAULT_PLATFORMS.keys())}")
            sys.exit(1)

    # Run benchmark
    engine = BenchmarkEngine(platforms, args.prompts_file)
    engine.run(max_prompts=args.prompts, category=args.category)

    # Summarize
    summaries = engine.summarize()

    # Print to terminal
    if not args.quiet:
        print_summary(summaries)

    # Generate report
    report = generate_markdown_report(summaries, engine.results, args.output)

    if not args.quiet:
        print(f"  Report saved to: {args.output}")
        print()

    return 0


if __name__ == "__main__":
    sys.exit(main())