sample.py

import argparse
import json
from pathlib import Path
from typing import List

from tqdm import tqdm

from api import create_api
from benchmark import create_benchmark


def generate_images(api_type: str, benchmarks: List[str]):
    images_dir = Path("images")
    api = create_api(api_type)

    api_dir = images_dir / api_type
    api_dir.mkdir(parents=True, exist_ok=True)

    for benchmark_type in tqdm(benchmarks, desc="Processing benchmarks"):
        print(f"\nProcessing benchmark: {benchmark_type}")

        benchmark = create_benchmark(benchmark_type)

        if benchmark_type == "geneval":
            benchmark_dir = api_dir / benchmark_type
            benchmark_dir.mkdir(parents=True, exist_ok=True)

            metadata_file = benchmark_dir / "metadata.jsonl"
            existing_metadata = {}
            if metadata_file.exists():
                with open(metadata_file, "r") as f:
                    for line in f:
                        entry = json.loads(line)
                        existing_metadata[entry["filepath"]] = entry

            with open(metadata_file, "a") as f:
                for metadata, folder_name in tqdm(
                    benchmark,
                    desc=f"Generating images for {benchmark_type}",
                    leave=False,
                ):
                    sample_path = benchmark_dir / folder_name
                    samples_path = sample_path / "samples"
                    samples_path.mkdir(parents=True, exist_ok=True)
                    image_path = samples_path / "0000.png"

                    if image_path.exists():
                        continue

                    try:
                        inference_time = api.generate_image(
                            metadata["prompt"], image_path
                        )

                        metadata_entry = {
                            "filepath": str(image_path),
                            "prompt": metadata["prompt"],
                            "inference_time": inference_time,
                        }

                        f.write(json.dumps(metadata_entry) + "\n")

                    except Exception as e:
                        print(
                            f"\nError generating image for prompt: {metadata['prompt']}"
                        )
                        print(f"Error: {str(e)}")
                        continue
        else:
            benchmark_dir = api_dir / benchmark_type
            benchmark_dir.mkdir(parents=True, exist_ok=True)

            metadata_file = benchmark_dir / "metadata.jsonl"
            existing_metadata = {}
            if metadata_file.exists():
                with open(metadata_file, "r") as f:
                    for line in f:
                        entry = json.loads(line)
                        existing_metadata[entry["filepath"]] = entry

            with open(metadata_file, "a") as f:
                for prompt, image_path in tqdm(
                    benchmark,
                    desc=f"Generating images for {benchmark_type}",
                    leave=False,
                ):
                    if image_path in existing_metadata:
                        continue

                    full_image_path = benchmark_dir / image_path

                    if full_image_path.exists():
                        continue

                    try:
                        inference_time = api.generate_image(prompt, full_image_path)

                        metadata_entry = {
                            "filepath": str(image_path),
                            "prompt": prompt,
                            "inference_time": inference_time,
                        }

                        f.write(json.dumps(metadata_entry) + "\n")

                    except Exception as e:
                        print(f"\nError generating image for prompt: {prompt}")
                        print(f"Error: {str(e)}")
                        continue


def main():
    parser = argparse.ArgumentParser(
        description="Generate images for specified benchmarks using a given API"
    )
    parser.add_argument("api_type", help="Type of API to use for image generation")
    parser.add_argument("benchmarks", nargs="+", help="List of benchmark types to run")

    args = parser.parse_args()

    generate_images(args.api_type, args.benchmarks)


if __name__ == "__main__":
    main()