File size: 4,850 Bytes
ffc54ec | 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 | """Paired held-out evaluation: 16 disjoint generation prompts + 2 disjoint edits, matched seeds.
precision: bf16 | fp8 | nvfp4. --accelerate uses the compiled runtime. BF16 must run first (edits
use BF16 outputs as their reference inputs for every precision). Saves PNGs and final latents.
"""
import argparse, json, time, collections
from pathlib import Path
import torch
from PIL import Image
from safetensors.torch import save_file
from common import ROOT, BASE, load_bf16_pipeline, sizes
from prompts import EVALUATION
EDITS = ['Replace the background with a blooming spring garden and preserve the animal.',
'Turn this room into a warm evening scene with lamps switched on, preserving its furniture.']
def load(precision, quant, offload=False):
if precision == 'bf16':
if offload:
from diffusers import QwenImage21Pipeline
pipe = QwenImage21Pipeline.from_pretrained(BASE, torch_dtype=torch.bfloat16)
pipe.set_progress_bar_config(disable=True)
return pipe
return load_bf16_pipeline()
if precision == 'fp8':
from fp8_runtime import load_pipeline
else:
from nvfp4_runtime import load_pipeline
return load_pipeline(str(BASE), quant)
@torch.inference_mode()
def main():
ap = argparse.ArgumentParser()
ap.add_argument('precision', choices=['bf16', 'fp8', 'nvfp4'])
ap.add_argument('--quant')
ap.add_argument('--name')
ap.add_argument('--accelerate', action='store_true')
ap.add_argument('--attention', default='bf16', choices=['bf16', 'sage'])
ap.add_argument('--fp8-first-steps', type=int, default=0)
ap.add_argument('--only', type=str, help='comma-separated subset of case ids, e.g. 00,03,edit-0')
ap.add_argument('--offload', action='store_true', help='model CPU offload (same arithmetic, lower peak VRAM; timings not meaningful)')
a = ap.parse_args()
name = a.name or a.precision
out = ROOT / 'evaluation' / name; out.mkdir(parents=True, exist_ok=True)
pipe = load(a.precision, a.quant or str(ROOT / a.precision / 'release' / 'transformer'), a.offload)
if a.offload:
pipe.to('cpu'); pipe.enable_model_cpu_offload()
if a.accelerate:
from acceleration import accelerate_pipeline
accelerate_pipeline(pipe, attention=a.attention)
if a.fp8_first_steps:
from nvfp4_runtime import enable_fp8_first_steps
enable_fp8_first_steps(pipe, str(ROOT / 'fp8' / 'release' / 'transformer'), a.fp8_first_steps,
accelerate=a.accelerate, attention=a.attention)
only = set(a.only.split(',')) if a.only else None
rows = []
for i, prompt in enumerate(EVALUATION):
key = f'{i:02d}'
if only and key not in only:
continue
w, h = sizes(i, 'evaluation')
latest = {}
def cb(p, step, t, kw):
if step == p._num_timesteps - 1:
latest['latents'] = kw['latents'].detach()
return kw
torch.cuda.synchronize(); start = time.perf_counter()
im = pipe(prompt=prompt, width=w, height=h, generator=torch.Generator('cuda').manual_seed(20000 + i),
callback_on_step_end=cb).images[0]
torch.cuda.synchronize(); seconds = time.perf_counter() - start
im.save(out / f'{key}.png')
save_file({'latents': latest['latents'].cpu().contiguous()}, str(out / f'{key}-latents.safetensors'))
rows.append({'id': key, 'prompt': prompt, 'seed': 20000 + i, 'width': w, 'height': h,
'steps': pipe._num_timesteps, 'seconds_including_first_use': seconds})
print(json.dumps({'event': 'evaluation', 'name': name, 'id': key, 'seconds': seconds}), flush=True)
for j, prompt in enumerate(EDITS):
key = f'edit-{j}'
if only and key not in only:
continue
idx = [0, 3][j]
ref = Image.open(ROOT / 'evaluation' / 'bf16' / f'{idx:02d}.png').resize((1024, 1024))
im = pipe(prompt=prompt, image=ref, width=1024, height=1024,
generator=torch.Generator('cuda').manual_seed(21000 + j)).images[0]
im.save(out / f'{key}.png')
rows.append({'id': key, 'prompt': prompt, 'seed': 21000 + j, 'reference': f'bf16/{idx:02d}.png'})
(out / 'runs.json').write_text(json.dumps({'precision': a.precision, 'accelerated': a.accelerate, 'quant': a.quant,
'torch': torch.__version__, 'gpu': torch.cuda.get_device_name(),
'transformer_dtypes': dict(collections.Counter(str(t.dtype) for t in pipe.transformer.state_dict().values())),
'rows': rows}, indent=2, ensure_ascii=False))
print('EVALUATION_COMPLETE', name, flush=True)
if __name__ == '__main__':
main()
|