"""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): from scripts.runtime import check_sampling_support check_sampling_support(QwenImage21Pipeline.load_config(str(model), local_files_only=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') from scripts.runtime import release_prefix_cache_before_decode release_prefix_cache_before_decode(pipe) 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() sample_sigmas = pipe.config.get('sample_sigmas') if sample_sigmas is not None and len(sample_sigmas) != args.steps: raise ValueError('Requested steps differ from the checkpoint sampling schedule') 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), sample_sigmas=sample_sigmas, prefix_cache_release=bool(getattr(pipe, '_prefix_cache_release', False)), pipeline_sha256=sha256(__import__('inspect').getfile(type(pipe)))) (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 case.get('seeds', 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 actual_steps = len(pipe.scheduler.timesteps) if actual_steps != args.steps: raise ValueError(f'Actual denoising steps {actual_steps} differ from requested {args.steps}') 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=actual_steps, sample_sigmas=sample_sigmas, scheduler_sigmas=pipe.scheduler.sigmas.cpu().tolist(), released_prefix_tensor_bytes=pipe._prefix_cache_telemetry['tensor_bytes'], 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', newline='\n') 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()