File size: 6,430 Bytes
ba48d54 | 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 | """Fresh-process memory residency and editing reproduction audit."""
import argparse
import gc
import json
import os
import subprocess
import time
from pathlib import Path
import torch
from PIL import Image
from scripts.benchmark import load_pipeline, generate, environment
from scripts.integrity import sha256
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--model', required=True)
ap.add_argument('--output', required=True)
ap.add_argument('--fix-offload', action='store_true')
ap.add_argument('--native-edit', action='store_true')
ap.add_argument('--vae-tiling', action='store_true')
ap.add_argument('--edit-only', action='store_true')
args = ap.parse_args()
out = Path(args.output)
out.mkdir(parents=True, exist_ok=False)
pipe = None
def snapshot(label):
torch.cuda.synchronize()
row = dict(label=label, pid=os.getpid(), allocated=torch.cuda.memory_allocated(),
reserved=torch.cuda.memory_reserved(), peak=torch.cuda.max_memory_allocated(),
free_total=list(torch.cuda.mem_get_info()))
if pipe is not None:
row['components'] = {}
for name in ['text_encoder', 'transformer', 'vae']:
component = getattr(pipe, name)
counts = {}
states = {}
for p in component.parameters():
key = f'{p.device}/{p.dtype}'
counts[key] = counts.get(key, 0) + p.numel()*p.element_size()
for attr in ('CB', 'SCB'):
v = getattr(p, attr, None)
if isinstance(v, torch.Tensor):
key = f'weight.{attr}/{v.device}/{v.dtype}'
states[key] = states.get(key, 0) + v.numel()*v.element_size()
for m in component.modules():
state = getattr(m, 'state', None)
if state:
for k, v in vars(state).items():
if isinstance(v, torch.Tensor):
key = f'{k}/{v.device}/{v.dtype}'
states[key] = states.get(key, 0) + v.numel()*v.element_size()
row['components'][name] = {'parameters': counts, 'bnb_state': states}
with (out/'memory.jsonl').open('a', encoding='utf-8') as f:
f.write(json.dumps(row)+'\n')
print(label, 'allocated GiB', round(row['allocated']/2**30, 3), flush=True)
(out/'environment.json').write_text(json.dumps(environment(), indent=2), encoding='utf-8')
(out/'gpu-before.txt').write_text(subprocess.check_output(['nvidia-smi'], text=True), encoding='utf-8')
snapshot('before_load')
pipe = load_pipeline(args.model, offload_fix=False)
snapshot('after_load_and_offload')
if args.fix_offload:
from scripts.runtime import enable_int8_cpu_offload
enable_int8_cpu_offload(pipe)
snapshot('after_fix')
if args.vae_tiling:
pipe.vae.enable_tiling()
# Wrapping pipeline boundaries keeps model hooks intact and avoids per-step profiling overhead.
for method in ['encode_prompt', 'prepare_latents']:
original = getattr(pipe, method)
def wrapper(*a, _fn=original, _name=method, **kw):
snapshot('before_'+_name)
result = _fn(*a, **kw)
snapshot('after_'+_name)
return result
setattr(pipe, method, wrapper)
orig_decode = pipe.vae.decode
def decode(*a, **kw):
snapshot('before_decode')
result = orig_decode(*a, **kw)
snapshot('after_decode')
return result
pipe.vae.decode = decode
cases = json.loads(Path('benchmarks/cases.json').read_text(encoding='utf-8'))
portrait, edit = cases[0], cases[-1]
runs = [('edit_fresh', edit, 42, None, True),
('portrait', portrait, 42, None, True),
('edit_repeat', edit, 42, None, True),
('edit_rgb', edit, 42, 'RGB', True),
('edit_no_cache', edit, 42, None, False)]
if args.native_edit:
runs = [('edit_native', edit, 42, None, True)]
elif args.edit_only:
runs = [('edit_fresh', edit, 42, None, True)]
for name, case, seed, mode, cache in runs:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
snapshot(name+'_start')
size = 2048 if args.native_edit else 1024
kwargs = dict(prompt=case['prompt'], width=size, height=size, output_resolution=size,
num_inference_steps=40, true_cfg_scale=1.0, use_kv_cache=cache,
generator=torch.Generator('cpu').manual_seed(seed))
if case.get('input'):
with Image.open(case['input']) as im:
kwargs['image'] = im.convert(mode) if mode else im.copy()
started = time.perf_counter()
with torch.inference_mode():
image = pipe(**kwargs).images[0]
torch.cuda.synchronize()
elapsed = time.perf_counter()-started
snapshot(name+'_end')
image.save(out/(name+'.png'))
row = dict(name=name, seed=seed, input_mode=mode or 'original', kv_cache=cache,
width=size, height=size, steps=40, cfg=1.0,
vae_tiling=args.vae_tiling,
seconds=elapsed, peak_allocated_bytes=torch.cuda.max_memory_allocated(),
image_sha256=sha256(out/(name+'.png')),
input_sha256=sha256(case['input']) if case.get('input') else None)
with (out/'records.jsonl').open('a', encoding='utf-8') as f:
f.write(json.dumps(row)+'\n')
if not args.native_edit and not args.edit_only:
# Isolate the common VAE from denoising and quantization.
with Image.open(edit['input']) as im:
pixels = pipe.image_processor.preprocess(im.copy()).unsqueeze(2)
with torch.inference_mode():
latents = pipe.vae.encode(pixels.to('cuda', dtype=pipe.vae.dtype)).latent_dist.mode()
reconstruction = pipe.vae.decode(latents, return_dict=False)[0][:, :, 0]
pipe.image_processor.postprocess(reconstruction, output_type='pil')[0].save(out/'vae_roundtrip.png')
# Instrumentation retains bound methods; process exit is the isolation boundary.
snapshot('before_process_exit')
if __name__ == '__main__':
main()
|