File size: 7,492 Bytes
6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a ba48d54 6f8c80a | 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 | """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):
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')
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()
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))
(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 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
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=args.steps,
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') 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()
|