File size: 8,644 Bytes
a5e8975 | 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 | """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()
|