File size: 8,460 Bytes
9116984 | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """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()
|