scripts/benchmark.py

"""Paired generation. One process loads one saved pipeline."""
import argparse
import gc
import importlib.metadata
import json
import os
import platform
import subprocess
import time
from pathlib import Path

import psutil
import torch
from PIL import Image

from scripts.editing_protocol import load_input, validate_seeds
from scripts.integrity import sha256
from scripts.provenance import model_identity
from scripts.runtime import load_pipeline


def apply_memory_cap(gib):
    if gib is None:
        return None
    if not torch.cuda.is_available():
        raise RuntimeError('A CUDA allocator cap requires CUDA')
    total = torch.cuda.get_device_properties(0).total_memory
    cap = int(gib * 2**30)
    if cap >= total:
        raise ValueError('Memory cap must be smaller than the physical GPU')
    torch.cuda.set_per_process_memory_fraction(cap / total)
    return {'cap_bytes': cap, 'cap_gib': gib, 'total_bytes': total}


def environment():
    packages = ['torch', 'diffusers', 'transformers', 'accelerate', 'sdnq',
                'safetensors', 'huggingface-hub', 'tokenizers', 'numpy', 'pillow', 'psutil']
    versions = {}
    for name in packages:
        try:
            versions[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            versions[name] = None
    if any(versions[name] is None for name in packages):
        missing = [name for name, version in versions.items() if version is None]
        raise RuntimeError(f'Missing required packages: {missing}')
    return {'python': platform.python_version(), 'platform': platform.platform(),
            'gpu': torch.cuda.get_device_name() if torch.cuda.is_available() else platform.processor(),
            'cuda': torch.version.cuda, 'mps': bool(torch.backends.mps.is_available()),
            'system_ram_bytes': psutil.virtual_memory().total,
            'packages': versions}


def memory_status():
    if not torch.cuda.is_available():
        return dict(allocated_bytes=0, reserved_bytes=0, cuda_free_bytes=0, cuda_total_bytes=0)
    torch.cuda.synchronize()
    free, total = torch.cuda.mem_get_info()
    return dict(allocated_bytes=torch.cuda.memory_allocated(),
                reserved_bytes=torch.cuda.memory_reserved(),
                cuda_free_bytes=free, cuda_total_bytes=total)


def case_seeds(case, default):
    seeds = [int(seed) for seed in case['seeds']] if 'seeds' in case else list(default)
    validate_seeds([case], seeds)
    return seeds


def generate(pipe, case, seed, width, height, steps):
    kwargs = dict(prompt=case['prompt'], width=width, height=height,
                  output_resolution=width, num_inference_steps=steps,
                  true_cfg_scale=1.0, use_kv_cache=True,
                  generator=torch.Generator('cpu').manual_seed(seed))
    if case.get('input'):
        kwargs['image'] = load_input(case['input'])
    with torch.inference_mode():
        return pipe(**kwargs).images[0]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', required=True)
    parser.add_argument('--output', required=True)
    parser.add_argument('--cases', default='benchmarks/cases.json')
    parser.add_argument('--case', default=None)
    parser.add_argument('--seeds', default='42,123')
    parser.add_argument('--width', type=int, default=1024)
    parser.add_argument('--height', type=int, default=1024)
    parser.add_argument('--steps', type=int, default=40)
    parser.add_argument('--offload', choices=['auto', 'model', 'group', 'resident'], default='auto')
    parser.add_argument('--memory-cap-gib', type=float)
    parser.add_argument('--no-warmup', action='store_true')
    parser.add_argument('--device', default=None)
    args = parser.parse_args()
    if min(args.width, args.height, args.steps) <= 0 or args.width % 32 or args.height % 32:
        parser.error('Use positive steps and dimensions divisible by 32')
    cap = apply_memory_cap(args.memory_cap_gib)
    cases = json.loads(Path(args.cases).read_text(encoding='utf-8'))
    if args.case:
        cases = [case for case in cases if case['id'] == args.case]
    if not cases:
        parser.error('No matching cases')
    default_seeds = [int(seed) for seed in args.seeds.split(',')]
    planned = [(case, seed) for case in cases for seed in case_seeds(case, default_seeds)]
    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=True)
    records_path = out / 'records.jsonl'
    if records_path.exists():
        raise FileExistsError(f'Use a fresh output directory: {out}')
    runtime = environment()
    runtime.update(pid=os.getpid(), before_load_memory=memory_status(),
                   memory_cap=cap, benchmark_sha256=sha256(__file__),
                   runtime_helper_sha256=sha256(Path(__file__).with_name('runtime.py')),
                   device_helper_sha256=sha256(Path(__file__).with_name('device.py')))
    if torch.cuda.is_available():
        for key, query in [('gpu_before', '--query-gpu=name,driver_version,memory.used,memory.free'),
                           ('gpu_processes_before', '--query-compute-apps=pid,process_name,used_memory')]:
            runtime[key] = subprocess.check_output(['nvidia-smi', query, '--format=csv'], text=True)
    runtime['model_identity'] = model_identity(args.model)
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()
    started = time.perf_counter()
    pipe = load_pipeline(args.model, offload=args.offload, device=args.device, local_files_only=True)
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    runtime.update(model=str(Path(args.model).resolve()),
                   load_seconds=time.perf_counter() - started,
                   load_peak_allocated_bytes=torch.cuda.max_memory_allocated() if torch.cuda.is_available() else 0,
                   offload=pipe.image21_runtime['offload'],
                   requested_offload=args.offload,
                   runtime_policy=pipe.image21_runtime,
                   warmup=not args.no_warmup,
                   after_load_memory=memory_status(),
                   generator_device='cpu', cases_sha256=sha256(args.cases))
    (out / 'environment.json').write_text(json.dumps(runtime, indent=2), encoding='utf-8')
    if not args.no_warmup:
        warmup_case, warmup_seed = planned[0]
        print('Full-settings warmup (excluded from measurements)', flush=True)
        generate(pipe, warmup_case, 0, args.width, args.height, args.steps)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
    for case, seed in planned:
        name = f'{case["id"]}-s{seed}'
        print(f'Generating {name}', flush=True)
        input_hash = sha256(case['input']) if case.get('input') else None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()
            torch.cuda.synchronize()
        before = memory_status()
        started = time.perf_counter()
        image = generate(pipe, case, seed, args.width, args.height, args.steps)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        elapsed = time.perf_counter() - started
        image_path = out / f'{name}.png'
        image.save(image_path)
        row = dict(case_id=case['id'], category=case['category'], prompt=case['prompt'],
                   seed=seed, source_seed=case.get('source_seed'),
                   width=args.width, height=args.height, steps=args.steps,
                   cfg=1.0, kv_cache=True, offload=runtime['offload'], input_sha256=input_hash,
                   before_memory=before, after_memory=memory_status(),
                   seconds=elapsed,
                   peak_allocated_bytes=torch.cuda.max_memory_allocated() if torch.cuda.is_available() else 0,
                   peak_reserved_bytes=torch.cuda.max_memory_reserved() if torch.cuda.is_available() else 0,
                   process_rss_after_bytes=psutil.Process().memory_info().rss,
                   image=image_path.name, image_sha256=sha256(image_path),
                   image_mode=image.mode, actual_size=list(image.size))
        with records_path.open('a', encoding='utf-8') as handle:
            handle.write(json.dumps(row, ensure_ascii=False) + '\n')
        print(f'{name}: {elapsed:.2f}s, allocated peak {row["peak_allocated_bytes"]/2**30:.2f} GiB', flush=True)


if __name__ == '__main__':
    main()