| """Run real image-conditioned decisions with the released JEVision checkpoint.""" |
| import argparse |
| import json |
| from pathlib import Path |
| import statistics |
| import time |
| import torch |
| from PIL import Image |
| from training.judge.model import ImageJudge |
| from training.judge.fast_engine import FastEngine, convert_projections |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument('--root', default='./data', help='Base model cache root') |
| p.add_argument('--checkpoint', default='.') |
| p.add_argument('--image', help='Optional PNG, JPEG or WebP image') |
| p.add_argument('--question', default='What is the dominant color in the image?') |
| p.add_argument('--choices', default='red,blue,green') |
| p.add_argument('--output', default='verification.json') |
| a = p.parse_args() |
| torch.set_num_threads(4) |
| started = time.perf_counter() |
| model = ImageJudge(Path(a.root), checkpoint=Path(a.checkpoint)).eval() |
| engine = FastEngine(model, resident=False, row_buckets=(1, 2, 3, 4), max_call_tokens=2560) |
| convert_projections(engine, 'w8a8', stage_on_host=True) |
| engine.release_original(cpu_embeddings=True) |
| engine.load_costs(Path(__file__).parent / 'training/judge/engine_costs_rtx5060ti_w8a8.json') |
| engine.prepare(graphs=False) |
| engine.enable_static_vision() |
| setup = time.perf_counter() - started |
| q = {'instructions': a.question, 'criteria': {x: x for x in a.choices.split(',')}} |
| images = [(Path(a.image).name, Image.open(a.image).convert('RGB'))] if a.image else [(color, Image.new('RGB', (224, 224), color)) for color in ('red','blue','green')] |
| rows = [] |
| for label, im in images: |
| engine.predict_many('Use the supplied image.', [q], image=im) |
| for repeat in range(3): |
| torch.cuda.synchronize(); begin = time.perf_counter() |
| r = engine.predict_many('Use the supplied image.', [q], image=im) |
| torch.cuda.synchronize(); elapsed = (time.perf_counter()-begin)*1000 |
| probs = r['results'][0]['probabilities'] |
| rows.append({'image': label, 'prediction': max(probs, key=probs.get), 'probabilities': probs, 'wall_ms': elapsed}) |
| report = {'gpu': torch.cuda.get_device_name(), 'setup_seconds': setup, 'cases': rows, 'median_ms': statistics.median(x['wall_ms'] for x in rows), 'scope': 'Synthetic color image smoke test; warm inference excludes startup, network and queue time.'} |
| Path(a.output).write_text(json.dumps(report, indent=2)) |
| print(json.dumps(report, indent=2)) |
| if not a.image: |
| assert all(x['image'] == x['prediction'] for x in rows), 'Image-conditioned smoke test failed' |
|
|
