BAAI/Brainmu-Spike

1#!/usr/bin/env python32import csv,json,os,threading,time,traceback,uuid,sys3from pathlib import Path4from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer5from urllib.parse import urlparse,parse_qs6 7ROOT=Path(__file__).resolve().parents[1]; PROJECT=ROOT8sys.path.insert(0,str(ROOT/'code'))9from project_config import load_config10SETTINGS=load_config()['generation']['inference']11for key,sub in {'HF_HOME':'huggingface','TORCH_HOME':'torch','TRITON_CACHE_DIR':'triton','TORCHINDUCTOR_CACHE_DIR':'torchinductor','TMPDIR':'tmp'}.items():12 os.environ[key]=str(ROOT/'.cache'/sub);Path(os.environ[key]).mkdir(parents=True,exist_ok=True)13os.environ.update(HF_HUB_OFFLINE='1',TRANSFORMERS_OFFLINE='1',TOKENIZERS_PARALLELISM='false')14DEFAULT_GPU_COUNT=815DEFAULTS=dict(model_path=str(ROOT/'checkpoints/Brainmu'),adapter=str(PROJECT/'checkpoints/lora/adapter.safetensors'),frontend=str(PROJECT.parent/'model.safetensors'),spike_dir=str(ROOT/'data/test/spike'),gt_dir=str(ROOT/'data/test/gt'))16lock=threading.RLock(); stop=threading.Event()17state=dict(phase='idle',message='请先加载模型，再加载测试数据。',model_loaded=False,dataset_count=0,requested=0,rows=[],error='',run_dir='',started=None,defaults=DEFAULTS)18engine=None; dataset=[]19 20def snapshot():21 with lock:22  d=dict(state);d['rows']=list(state['rows'])23  d['prompt_supported']=True24  d['inference_settings']=SETTINGS25  d['parallel_supported']=True26  d['default_gpu_count']=DEFAULT_GPU_COUNT27  d['completed']=len(d['rows'])28  d['mean_psnr']=sum(x['psnr_db'] for x in d['rows'])/len(d['rows']) if d['rows'] else None29  d['mean_ssim']=sum(x['ssim'] for x in d['rows'])/len(d['rows']) if d['rows'] else None30  d['full_1000_complete']=d.get('evaluation_complete',d['phase']=='complete') and d['completed']==d['requested']==100031  d['elapsed']=round(time.time()-d['started'],1) if d['started'] and d['phase']=='running' else state.get('elapsed',0)32  return d33 34def persist():35 d=snapshot();out=Path(d['run_dir'])36 tmp=out/'metrics.json.tmp';tmp.write_text(json.dumps(d,ensure_ascii=False,indent=2));tmp.replace(out/'metrics.json')37 fields=['id','index','seed','prompt','gpu_id','psnr_db','ssim','spike','gt','condition','prediction']38 tmp=out/'metrics.csv.tmp'39 with tmp.open('w',newline='') as f:40  w=csv.DictWriter(f,fieldnames=fields);w.writeheader();w.writerows(sorted(d['rows'],key=lambda r:r['index']))41 tmp.replace(out/'metrics.csv')42 43def job(fn, persist_run=False):44 def wrapped():45  try:fn()46  except Exception as e:47   traceback.print_exc()48   with lock:state.update(phase='error',error=str(e),message='操作失败，请查看错误信息。')49   if persist_run:50    with lock:state['model_loaded']=False51    if state['run_dir']:persist()52 threading.Thread(target=wrapped,daemon=True).start()53 54def action(kind,data):55 global engine,dataset56 with lock:57  busy=state['phase'] in ['loading','validating','running']58  if kind=='stop':59   stop.set();state['message']='将在当前样本完成后停止。';return60  if busy:raise ValueError('任务正在执行，请稍候。')61  if kind=='load':62   if state['model_loaded']:return63   config={k:str(Path(data.get(k,DEFAULTS[k])).expanduser().resolve()) for k in ['model_path','adapter','frontend']}64   for p in [Path(config['model_path'])/n for n in ['ema.safetensors','ae.safetensors','llm_config.json','vit_config.json']]+[Path(config['adapter']),Path(config['adapter']).with_name('adapter_config.json'),Path(config['frontend'])]:65    if not p.is_file():raise ValueError('文件不存在：'+str(p))66   gpu_count=int(data.get('gpu_count',DEFAULT_GPU_COUNT))67   if gpu_count not in [1,2,4,8]:raise ValueError('GPU 数量请选择 1、2、4 或 8。')68   state.update(phase='loading',gpu_count=gpu_count,gpu_ids=list(range(gpu_count)),message=f'正在向 {gpu_count} 张 GPU 加载 Brainmu 模型…',error='')69   def load():70    global engine71    from multi_gpu import MultiGPUEngine72    candidate=MultiGPUEngine(list(range(gpu_count)));candidate.load(**config);engine=candidate73    with lock:state.update(phase='ready',model_loaded=True,model_config=config,message=f'{gpu_count} 卡模型已加载，可以选择测试集。')74   job(load)75  elif kind=='dataset':76   if not state['model_loaded']:raise ValueError('请先加载模型。')77   sd=Path(data.get('spike_dir',DEFAULTS['spike_dir'])).expanduser().resolve();gd=Path(data.get('gt_dir',DEFAULTS['gt_dir'])).expanduser().resolve()78   dataset=[]79   state.update(phase='validating',evaluation_complete=False,dataset_count=0,rows=[],requested=0,run_dir='',started=None,elapsed=0,message='检查 1,000 对 DAT 与 GT…',error='')80   def validate():81    global dataset82    from PIL import Image83    ps=sorted(sd.glob('*.dat'))84    if len(ps)!=1000:raise ValueError(f'当前目录有 {len(ps)} 个 DAT，要求恰好 1,000 个。')85    pairs=[]86    for p in ps:87     gt=gd/(p.stem+'.png')88     if p.stat().st_size!=512500:raise ValueError('DAT 应为 41×250×400 位：'+p.name)89     if not gt.is_file():raise ValueError('缺少对应 GT：'+gt.name)90     with Image.open(gt) as im:91      if im.size!=(400,250):raise ValueError('GT 应为 400×250：'+gt.name)92      im.verify()93     pairs.append((p,gt))94    with lock:dataset=pairs;state.update(phase='ready',dataset_count=len(pairs),dataset_config=dict(spike_dir=str(sd),gt_dir=str(gd)),message='已加载 1,000 对 DAT / GT，可以开始测试。')95   job(validate)96  elif kind=='run':97   if not state['model_loaded'] or len(dataset)!=1000:raise ValueError('请先加载模型和 1,000 对数据。')98   limit=int(data.get('limit',1000))99   prompt=data.get('prompt',SETTINGS['prompt'])100   if not isinstance(prompt,str) or not prompt.strip() or len(prompt)>2000:raise ValueError('Prompt 必须为 1–2000 个字符。')101   prompt=prompt.strip()102   if limit not in [3,16,1000]:raise ValueError('请选择 3 张、16 张或完整 1,000 张。')103   out=ROOT/'outputs/ui_runs'/(time.strftime('%Y%m%d_%H%M%S')+'_'+uuid.uuid4().hex[:6])104   for sub in ['condition','prediction']:(out/sub).mkdir(parents=True)105   stop.clear();state.update(phase='running',evaluation_complete=False,run_prompt=prompt,requested=limit,rows=[],run_dir=str(out),started=time.time(),elapsed=0,error='',message='正在重构和计算逐图指标…')106   pairs=list(dataset[:limit])107   def run():108    persist()109    for row in engine.predict_many(pairs,out,prompt,stop):110     with lock:111      state['rows'].append(row)112      state['message']=f"已完成 {len(state['rows'])} / {limit}：{row['id']}（GPU {row['gpu_id']}）"113     persist()114    with lock:115     complete=len(state['rows'])==limit116     state.update(phase='complete' if complete else 'stopped',evaluation_complete=complete,message='测试完成。' if complete else '已停止，结果已保存。',elapsed=round(time.time()-state['started'],1))117    persist()118   job(run,persist_run=True)119  else:raise ValueError('未知操作')120 121class Handler(BaseHTTPRequestHandler):122 def send(self,body,ctype='application/json',status=200,extra=None):123  if isinstance(body,(dict,list)):body=json.dumps(body,ensure_ascii=False).encode()124  if isinstance(body,str):body=body.encode()125  self.send_response(status);self.send_header('Content-Type',ctype);self.send_header('Content-Length',str(len(body)));self.send_header('Cache-Control','no-store')126  self.send_header('X-Content-Type-Options','nosniff')127  for k,v in (extra or {}).items():self.send_header(k,v)128  self.end_headers();self.wfile.write(body)129 def do_GET(self):130  p=urlparse(self.path)131  try:132   if p.path=='/':self.send((Path(__file__).parent/'index.html').read_bytes(),'text/html; charset=utf-8')133   elif p.path=='/api/state':self.send(snapshot())134   elif p.path=='/image':135    q=parse_qs(p.query);i=int(q['i'][0]);kind=q['kind'][0]136    if kind not in ['condition','prediction','gt'] or i<0:raise ValueError('Invalid image')137    with lock:row=dict(state['rows'][i])138    self.send(Path(row[kind]).read_bytes(),'image/png')139   elif p.path=='/download':140    q=parse_qs(p.query);ext=q.get('format',['csv'])[0]141    if ext not in ['csv','json']:raise ValueError('Invalid format')142    with lock:run=state['run_dir']143    if not run:raise ValueError('尚无结果')144    self.send((Path(run)/('metrics.'+ext)).read_bytes(),'text/csv; charset=utf-8' if ext=='csv' else 'application/json',extra={'Content-Disposition':f'attachment; filename="metrics.{ext}"'})145   else:self.send({'error':'Not found'},status=404)146  except (ValueError,KeyError,IndexError,FileNotFoundError) as e:self.send({'error':str(e)},status=400)147 def do_POST(self):148  origin=self.headers.get('Origin')149  if origin and origin not in ['http://'+self.headers.get('Host','')]:150   self.send({'error':'Origin rejected'},status=403);return151  try:152   if self.headers.get_content_type()!='application/json':raise ValueError('JSON required')153   length=int(self.headers.get('Content-Length','0'))154   if not 0<length<16384:raise ValueError('Invalid body')155   action(urlparse(self.path).path.removeprefix('/api/'),json.loads(self.rfile.read(length)));self.send({'ok':True})156  except (ValueError,TypeError) as e:self.send({'error':str(e)},status=400)157 def log_message(self,fmt,*args):158  if '/api/state' not in (args[0] if args else ''):super().log_message(fmt,*args)159 160 161if __name__=='__main__':162 import argparse163 ap=argparse.ArgumentParser();ap.add_argument('--port',type=int,default=8997);args=ap.parse_args()164 print(f'Listening on http://127.0.0.1:{args.port}',flush=True)165 ThreadingHTTPServer(('127.0.0.1',args.port),Handler).serve_forever()166