BAAI/Brainmu-Spike

1"""Brainmu: reproducible DAT reconstruction and grayscale evaluation."""2import argparse,csv,hashlib,json,math,platform,sys,time,traceback3from pathlib import Path4ROOT=Path(__file__).resolve().parents[1]5sys.path.insert(0,str(ROOT/'code'))6from project_config import load_config,DEFAULT_CONFIG7 8def sha256(path):9 h=hashlib.sha256()10 with path.open('rb') as f:11  for block in iter(lambda:f.read(1024*1024),b''):h.update(block)12 return h.hexdigest()13 14def validate_pairs(spike_dir,gt_dir,expected_count=0):15 from PIL import Image16 spec=load_config()['input']17 files=sorted(spike_dir.glob('*.dat'))18 if not files:raise ValueError('No DAT files found')19 if expected_count and len(files)!=expected_count:raise ValueError(f'Expected {expected_count} DAT files, found {len(files)}')20 rows=[]21 for f in files:22  gt=gt_dir/(f.stem+'.png')23  if f.stat().st_size!=spec['packed_bytes']:raise ValueError('DAT size differs from config.json: '+f.name)24  if not gt.is_file():raise ValueError('Missing GT: '+gt.name)25  with Image.open(gt) as im:26   if im.size!=(spec['width'],spec['height']):raise ValueError('GT must be 400 x 250: '+gt.name)27   im.verify()28  rows.append(dict(id=f.stem,spike=str(f),gt=str(gt),spike_sha256=sha256(f),gt_sha256=sha256(gt)))29 return files,rows30 31def main():32 project=load_config();inference=project['generation']['inference']33 p=argparse.ArgumentParser(description=__doc__)34 for name,default in [('model-path','checkpoints/Brainmu'),('adapter','checkpoints/lora/adapter.safetensors'),('frontend','../model.safetensors'),('spike-dir','data/test/spike'),('gt-dir','data/test/gt')]:35  p.add_argument('--'+name,type=Path,default=ROOT/default)36 p.add_argument('--output-dir',type=Path,default=ROOT/'outputs'/('test_'+time.strftime('%Y%m%d_%H%M%S')))37 p.add_argument('--expected-count',type=int,default=0,help='Require exactly this many pairs; 0 = any count')38 p.add_argument('--limit',type=int,default=0,help='Smoke test only: sorted prefix; incompatible with expected-count')39 p.add_argument('--prompt',default=inference['prompt'])40 p.add_argument('--check-only',action='store_true',help='Validate all inputs without loading models or using GPU')41 a=p.parse_args()42 for name in ['model_path','adapter','frontend','spike_dir','gt_dir','output_dir']:setattr(a,name,getattr(a,name).resolve())43 if a.limit<0 or a.expected_count<0:p.error('Counts must be nonnegative')44 if a.limit and a.expected_count:p.error('For subset smoke tests set --expected-count 0; full evaluation must not use --limit')45 if not a.prompt.strip():p.error('Prompt must not be empty')46 if a.output_dir.exists():p.error('output-dir must be new')47 required=[a.model_path/n for n in ['ema.safetensors','ae.safetensors','llm_config.json','vit_config.json','tokenizer_config.json','tokenizer.json']]+[a.adapter,a.adapter.with_name('adapter_config.json'),a.frontend]48 for f in required:49  if not f.is_file():p.error('Missing model file: '+str(f))50 try:files,manifest=validate_pairs(a.spike_dir,a.gt_dir,a.expected_count)51 except (ValueError,OSError) as exc:p.error(str(exc))52 if a.limit:files=files[:a.limit];manifest=manifest[:a.limit]53 a.output_dir.mkdir(parents=True)54 for sub in ['condition','prediction']:(a.output_dir/sub).mkdir()55 (a.output_dir/'dataset_manifest.json').write_text(json.dumps(manifest,indent=2))56 import importlib.metadata as metadata57 versions={}58 for name in ['torch','torchvision','transformers','accelerate','flash-attn','numpy','Pillow']:59  try:versions[name]=metadata.version(name)60  except metadata.PackageNotFoundError:versions[name]=None61 config=dict(arguments={k:str(v) if isinstance(v,Path) else v for k,v in vars(a).items()},python=platform.python_version(),packages=versions,project_config=project,config_sha256=sha256(DEFAULT_CONFIG),metric='Pillow RGB->L, [0,1], per-image PSNR then mean; SSIM Gaussian 11x11 sigma=1.5, no border crop; GT resized to prediction',weights=[dict(path=str(f),bytes=f.stat().st_size,mtime_ns=f.stat().st_mtime_ns) for f in required],weight_fingerprint='size/mtime only, not content hashes',source_sha256={str(f.relative_to(ROOT)):sha256(f) for base in ['code','ui','vendor/Brainmu'] for f in sorted((ROOT/base).rglob('*.py'))})62 (a.output_dir/'run_config.json').write_text(json.dumps(config,indent=2))63 rows=[];begin=time.time()64 def persist(status,error=None):65  n=len(rows);complete=status=='complete' and n==len(files)66  report=dict(status=status,completed=n,requested=len(files),complete=complete,full_evaluation=complete and not a.limit,expected_count=a.expected_count,mean_psnr=sum(x['psnr_db'] for x in rows)/n if n else None,mean_ssim=sum(x['ssim'] for x in rows)/n if n else None,elapsed_seconds=time.time()-begin,error=error,rows=rows)67  tmp=a.output_dir/'metrics.json.tmp';tmp.write_text(json.dumps(report,indent=2,allow_nan=False));tmp.replace(a.output_dir/'metrics.json')68  if rows:69   tmp=a.output_dir/'metrics.csv.tmp'70   with tmp.open('w',newline='') as h:71    w=csv.DictWriter(h,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)72   tmp.replace(a.output_dir/'metrics.csv')73  return report74 if a.check_only:75  persist('validated');print(f'VALIDATED {len(files)} pairs; no inference performed. Report: {a.output_dir}');return76 persist('loading')77 try:78  sys.path.insert(0,str(ROOT/'ui'))79  from engine import Engine80  e=Engine();e.load(a.model_path,a.adapter,a.frontend)81  for i,f in enumerate(files):82   row=e.predict(f,a.gt_dir/(f.stem+'.png'),i,a.output_dir,a.prompt)83   if not all(math.isfinite(row[k]) for k in ['psnr_db','ssim']):raise ValueError('Nonfinite metric: '+f.name)84   rows.append(row);persist('running')85   print(f"{len(rows)}/{len(files)} PSNR={row['psnr_db']:.4f} SSIM={row['ssim']:.6f}",flush=True)86  report=persist('complete')87  print(f"COMPLETE {len(rows)}/{len(files)} | mean PSNR {report['mean_psnr']:.6f} dB | mean SSIM {report['mean_ssim']:.6f} | {a.output_dir}",flush=True)88 except BaseException as exc:89  persist('interrupted' if isinstance(exc,KeyboardInterrupt) else 'failed',str(exc));raise90if __name__=='__main__':main()91