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):
    from scripts.runtime import check_sampling_support
    check_sampling_support(QwenImage21Pipeline.load_config(str(model), local_files_only=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')
    from scripts.runtime import release_prefix_cache_before_decode
    release_prefix_cache_before_decode(pipe)
    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()
    sample_sigmas = pipe.config.get('sample_sigmas')
    if sample_sigmas is not None and len(sample_sigmas) != args.steps:
        raise ValueError('Requested steps differ from the checkpoint sampling schedule')
    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),
                   sample_sigmas=sample_sigmas,
                   prefix_cache_release=bool(getattr(pipe, '_prefix_cache_release', False)),
                   pipeline_sha256=sha256(__import__('inspect').getfile(type(pipe))))
    (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 case.get('seeds', 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
            actual_steps = len(pipe.scheduler.timesteps)
            if actual_steps != args.steps:
                raise ValueError(f'Actual denoising steps {actual_steps} differ from requested {args.steps}')
            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=actual_steps,
                       sample_sigmas=sample_sigmas, scheduler_sigmas=pipe.scheduler.sigmas.cpu().tolist(),
                       released_prefix_tensor_bytes=pipe._prefix_cache_telemetry['tensor_bytes'],
                       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', newline='\n') 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()