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