"""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()