File size: 4,244 Bytes
1435032 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """Measure the declared 2048px editing suite in a fresh process per precision."""
import argparse
import gc
import json
import os
import subprocess
import time
from pathlib import Path
import psutil
import torch
from scripts.benchmark import environment, generate, load_pipeline, memory_status
from scripts.integrity import sha256
from scripts.provenance import model_identity
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--model', required=True)
ap.add_argument('--output', required=True)
ap.add_argument('--cases', default='benchmarks/editing.json')
ap.add_argument('--inputs', required=True, help='Directory containing the published input PNGs')
args = ap.parse_args()
cases = json.loads(Path(args.cases).read_text(encoding='utf-8'))
for case in cases:
case['input'] = str(Path(args.inputs) / case['input_file'])
if sha256(case['input']) != case['input_sha256']:
raise ValueError('Input hash mismatch: ' + case['id'])
out = Path(args.output)
out.mkdir(parents=True, exist_ok=False)
env = environment()
env.update(pid=os.getpid(), before_load_memory=memory_status(),
model_identity=model_identity(args.model), offload='model', offload_aux_fix=True,
warmup=True, generator_device='cpu', cases_sha256=sha256(args.cases),
source_sha256={n: sha256(Path(__file__).with_name(n)) for n in
('benchmark_edits.py', 'benchmark.py', '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')]:
env[key] = subprocess.check_output(['nvidia-smi', query, '--format=csv'], text=True)
if env['before_load_memory']['allocated_bytes'] != 0:
raise ValueError('Expected a fresh CUDA allocation baseline')
torch.cuda.reset_peak_memory_stats()
start = time.perf_counter()
pipe = load_pipeline(args.model)
pipe.vae.enable_tiling()
env.update(load_seconds=time.perf_counter()-start,
load_peak_allocated_bytes=torch.cuda.max_memory_allocated(),
after_load_memory=memory_status(), vae_tiling=True,
vae_tiles={n: getattr(pipe.vae, n) for n in
('tile_sample_min_height', 'tile_sample_min_width',
'tile_sample_stride_height', 'tile_sample_stride_width')})
(out/'environment.json').write_text(json.dumps(env, indent=2), encoding='utf-8')
print('Full-settings 2048px warmup, seed 0 (excluded)', flush=True)
generate(pipe, cases[0], 0, 2048, 2048, 40)
torch.cuda.synchronize()
for case in cases:
print('Generating ' + case['id'], flush=True)
gc.collect()
torch.cuda.empty_cache()
before = memory_status()
torch.cuda.reset_peak_memory_stats()
start = time.perf_counter()
image = generate(pipe, case, 42, 2048, 2048, 40)
torch.cuda.synchronize()
seconds = time.perf_counter()-start
path = out/(case['id']+'-s42.png')
image.save(path)
row = dict(case_id=case['id'], category=case['category'], prompt=case['prompt'],
seed=42, width=2048, height=2048, output_resolution=2048, steps=40,
cfg=1.0, kv_cache=True, offload='model', vae_tiling=True,
input_file=case['input_file'], input_sha256=case['input_sha256'],
before_memory=before, after_memory=memory_status(), seconds=seconds,
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=path.name, image_sha256=sha256(path), image_mode=image.mode,
actual_size=list(image.size))
with (out/'records.jsonl').open('a', encoding='utf-8') as stream:
stream.write(json.dumps(row, ensure_ascii=False)+'\n')
print(f'{case["id"]}: {seconds:.2f}s, {row["peak_allocated_bytes"]/2**30:.2f} GiB', flush=True)
if __name__ == '__main__':
main()
|