scripts/benchmark.py

"""Paired reproducible generation; each invocation 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 diffusers import QwenImage21Pipeline
from scripts.integrity import sha256
from scripts.provenance import model_identity


def environment():
    packages = ['torch', 'torchvision', 'diffusers', 'transformers', 'accelerate',
                'bitsandbytes', 'safetensors', 'huggingface-hub', 'tokenizers',
                'numpy', 'pillow', 'psutil']
    return {'python': platform.python_version(), 'platform': platform.platform(),
            'gpu': torch.cuda.get_device_name(), 'cuda': torch.version.cuda,
            'packages': {p: importlib.metadata.version(p) for p in packages}}


def memory_status():
    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 load_pipeline(model, offload='model', offload_fix=True):
    config = Path(model) / 'transformer/config.json'
    if offload == 'model' and offload_fix and config.exists():
        qconfig = json.loads(config.read_text()).get('quantization_config', {})
        if qconfig.get('load_in_8bit'):
            from scripts.runtime import load_int8_pipeline
            return load_int8_pipeline(str(model), local_files_only=True)
    pipe = QwenImage21Pipeline.from_pretrained(str(model), dtype=torch.bfloat16,
                                               local_files_only=True)
    if offload == 'model':
        if offload_fix:
            from scripts.runtime import enable_int8_cpu_offload
            enable_int8_cpu_offload(pipe)
        else:
            pipe.enable_model_cpu_offload()
    else:
        pipe.to('cuda')
    return pipe


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'):
        with Image.open(case['input']) as image:
            kwargs['image'] = image.copy()
    with torch.inference_mode():
        return pipe(**kwargs).images[0]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', required=True)
    ap.add_argument('--output', required=True)
    ap.add_argument('--cases', default='benchmarks/cases.json')
    ap.add_argument('--case', default=None)
    ap.add_argument('--seeds', default='42,123')
    ap.add_argument('--width', type=int, default=1024)
    ap.add_argument('--height', type=int, default=1024)
    ap.add_argument('--steps', type=int, default=40)
    ap.add_argument('--offload', choices=['model', 'none'], default='model')
    ap.add_argument('--no-warmup', action='store_true')
    ap.add_argument('--legacy-offload', action='store_true', help='Reproduce the original auxiliary-tensor retention bug')
    args = ap.parse_args()
    if min(args.width, args.height, args.steps) <= 0 or args.width % 32 or args.height % 32:
        ap.error('Use positive steps and dimensions divisible by 32')
    cases = json.loads(Path(args.cases).read_text(encoding='utf-8'))
    if args.case:
        cases = [c for c in cases if c['id'] == args.case]
    if not cases:
        ap.error('No matching cases')
    seeds = [int(s) for s in args.seeds.split(',')]
    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(),
                   offload_aux_fix=not args.legacy_offload,
                   benchmark_sha256=sha256(__file__),
                   runtime_helper_sha256=sha256(Path(__file__).with_name('runtime.py')))
    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)
    torch.cuda.reset_peak_memory_stats()
    started = time.perf_counter()
    pipe = load_pipeline(args.model, args.offload, not args.legacy_offload)
    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(),
                   offload=args.offload, 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:
        print('Full-settings warmup (excluded from measurements)', flush=True)
        generate(pipe, cases[0], 0, args.width, args.height, args.steps)
        torch.cuda.synchronize()
    for case in cases:
        for seed in seeds:
            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()
            torch.cuda.empty_cache()
            before = memory_status()
            torch.cuda.reset_peak_memory_stats()
            torch.cuda.synchronize()
            started = time.perf_counter()
            image = generate(pipe, case, seed, args.width, args.height, args.steps)
            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, width=args.width, height=args.height, steps=args.steps,
                       cfg=1.0, kv_cache=True, offload=args.offload, input_sha256=input_hash,
                       before_memory=before, after_memory=memory_status(),
                       seconds=elapsed, peak_allocated_bytes=torch.cuda.max_memory_allocated(),
                       peak_reserved_bytes=torch.cuda.max_memory_reserved(),
                       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))
            if image.mode == 'RGBA':
                import numpy as np
                alpha = np.asarray(image.getchannel('A'))
                row['alpha_min'] = int(alpha.min())
                row['alpha_max'] = int(alpha.max())
                row['alpha_nonopaque_fraction'] = float((alpha < 255).mean())
            with records_path.open('a', encoding='utf-8') as f:
                f.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()