| import base64,io,json,re,threading |
| from pathlib import Path |
| import numpy as np |
| from PIL import Image |
| import torch |
| from .networks import SketchEncoder,SpikingReadout,PongSNN,CATEGORIES,NavigationPolicy |
|
|
| ROOT=Path(__file__).resolve().parents[1] |
| DEVICE='mps' if torch.backends.mps.is_available() else 'cpu' |
| torch.set_num_threads(12) |
| LOCK=threading.RLock() |
| CACHE={} |
|
|
| def decode_image(data): |
| if not isinstance(data,str) or len(data)>12_000_000: raise ValueError('Image is too large.') |
| raw=base64.b64decode(data.split(',',1)[-1],validate=True) |
| im=Image.open(io.BytesIO(raw)) |
| if im.width*im.height>16_000_000: raise ValueError('Image exceeds 16 megapixels.') |
| return im.convert('RGB') |
|
|
| def encode_image(im): |
| b=io.BytesIO(); im.save(b,format='PNG') |
| return 'data:image/png;base64,'+base64.b64encode(b.getvalue()).decode() |
|
|
| def load(name): |
| if name in CACHE: return CACHE[name] |
| path=ROOT/'models'/'hf'/name |
| if not (path/'config.json').exists() and name!='image': |
| raise RuntimeError(f'{name} model has not been downloaded. Run scripts/download_models.py.') |
| from transformers import AutoTokenizer,AutoModelForCausalLM,CLIPProcessor,CLIPModel,BlipProcessor,BlipForConditionalGeneration |
| if name=='language': |
| tok=AutoTokenizer.from_pretrained(path,local_files_only=True) |
| model=AutoModelForCausalLM.from_pretrained(path,local_files_only=True,torch_dtype=torch.float32).to(DEVICE).eval() |
| obj=(tok,model) |
| elif name=='vision': |
| obj=(CLIPProcessor.from_pretrained(path,local_files_only=True),CLIPModel.from_pretrained(path,local_files_only=True).to(DEVICE).eval()) |
| elif name=='caption': |
| obj=(BlipProcessor.from_pretrained(path,local_files_only=True),BlipForConditionalGeneration.from_pretrained(path,local_files_only=True).to(DEVICE).eval()) |
| elif name=='speech': |
| from transformers import WhisperProcessor,WhisperForConditionalGeneration |
| obj=(WhisperProcessor.from_pretrained(path,local_files_only=True),WhisperForConditionalGeneration.from_pretrained(path,local_files_only=True).to(DEVICE).eval()) |
| elif name=='image': |
| from diffusers import AutoPipelineForText2Image |
| |
| |
| obj=AutoPipelineForText2Image.from_pretrained(path,local_files_only=True,torch_dtype=torch.float32,variant='fp16',use_safetensors=True).to(DEVICE) |
| obj.enable_attention_slicing() |
| else: raise ValueError('Unknown model') |
| CACHE[name]=obj |
| return obj |
|
|
| def language(prompt,max_tokens=160): |
| with LOCK,torch.inference_mode(): |
| tok,model=load('language') |
| text=tok.apply_chat_template([{'role':'system','content':'You are a helpful assistant for a classroom AI brain exhibit. Follow the requested format exactly. Be brief.'},{'role':'user','content':prompt}],tokenize=False,add_generation_prompt=True) |
| inputs=tok(text,return_tensors='pt').to(DEVICE) |
| out=model.generate(**inputs,max_new_tokens=max_tokens,do_sample=False,pad_token_id=tok.eos_token_id) |
| return tok.decode(out[0,inputs.input_ids.shape[1]:],skip_special_tokens=True).strip() |
|
|
| def image_features(im): |
| with LOCK,torch.inference_mode(): |
| proc,model=load('vision'); data=proc(images=im,return_tensors='pt').to(DEVICE) |
| feat=model.get_image_features(**data); feat=feat/feat.norm(dim=-1,keepdim=True) |
| return feat[0].cpu().numpy() |
|
|
| def text_features(texts): |
| with LOCK,torch.inference_mode(): |
| proc,model=load('vision'); data=proc(text=texts,return_tensors='pt',padding=True,truncation=True).to(DEVICE) |
| feat=model.get_text_features(**data); feat=feat/feat.norm(dim=-1,keepdim=True) |
| return feat.cpu().numpy() |
|
|
| def classify_image(im,labels): |
| feat=image_features(im); refs=text_features(['a photo of '+label for label in labels]) |
| scores=torch.tensor(refs@feat*30).softmax(0).numpy() |
| return sorted([{'label':l,'probability':float(p)} for l,p in zip(labels,scores)],key=lambda a:-a['probability']),feat |
|
|
| def caption(im): |
| with LOCK,torch.inference_mode(): |
| proc,model=load('caption'); data=proc(images=im,return_tensors='pt').to(DEVICE) |
| out=model.generate(**data,max_new_tokens=45) |
| return proc.decode(out[0],skip_special_tokens=True) |
|
|
| def generate_image(prompt,seed=21): |
| with LOCK,torch.inference_mode(): |
| pipe=load('image'); gen=torch.Generator(device='cpu').manual_seed(seed) |
| im=pipe(prompt=prompt,num_inference_steps=2,guidance_scale=0.,height=512,width=512,generator=gen).images[0] |
| return encode_image(im) |
|
|
| def sketch_input(im): |
| gray=255-np.asarray(im.convert('L')) |
| ys,xs=np.where(gray>30) |
| if len(xs)<8: return None |
| bbox=(int(xs.min()),int(ys.min()),int(xs.max())+1,int(ys.max())+1) |
| crop=Image.fromarray(gray).crop(bbox) |
| crop.thumbnail((25,25),Image.Resampling.LANCZOS) |
| out=Image.new('L',(28,28)); out.paste(crop,((28-crop.width)//2,(28-crop.height)//2)) |
| return torch.tensor(np.asarray(out,dtype=np.float32)[None,None]/255) |
|
|
| def sketch(im,threshold=1.): |
| x=sketch_input(im) |
| if x is None: return {'candidates':[],'spikes':[],'membrane':[],'message':'Add a few strokes to begin.'} |
| with LOCK,torch.inference_mode(): |
| if 'sketch' not in CACHE: |
| ck=torch.load(ROOT/'models'/'sketch.pt',map_location='cpu',weights_only=True) |
| enc=SketchEncoder(); enc.load_state_dict(ck['encoder']); enc.eval() |
| net=SpikingReadout(); net.load_state_dict(ck['snn']); net.eval() |
| CACHE['sketch']=(enc,net,ck['mean'],ck['std']) |
| enc,net,mean,std=CACHE['sketch'] |
| features=(enc.features(x)-mean)/std |
| logits,spikes,mem=net(features,threshold,True) |
| scores=logits[0].softmax(0).numpy() |
| ids=np.argsort(-scores)[:5] |
| return {'candidates':[{'label':CATEGORIES[i],'probability':float(scores[i])} for i in ids], |
| 'spikes':spikes[:,0,:32].int().tolist(),'membrane':mem[:,0,0].tolist(), |
| 'spike_count':int(spikes.sum()),'threshold':threshold,'model':'CNN encoder → trained LIF layer → linear class readout'} |
|
|
| def pong_config(): |
| receipt=json.loads((ROOT/'models'/'pong_v2_receipt.json').read_text()) |
| return {'model':'Target-position LIF SNN','optimal_threshold':receipt['optimal_threshold'], |
| 'steps':10,'neurons':64,'validation_target_mae':receipt['threshold_validation_mae'][str(receipt['optimal_threshold'])], |
| 'imitation_accuracy':receipt['imitation_accuracy'],'test_states':receipt['test_states'], |
| 'benchmark_scope':receipt['benchmark_scope']} |
|
|
| def pong(state,threshold): |
| if len(state)!=5 or not np.isfinite(state).all():raise ValueError('Use five finite game observations.') |
| with LOCK,torch.inference_mode(): |
| if 'pong' not in CACHE: |
| net=PongSNN();net.load_state_dict(torch.load(ROOT/'models'/'pong_v2.pt',weights_only=True));net.eval();CACHE['pong']=net |
| goal,spikes,mem=CACHE['pong'](torch.tensor([state],dtype=torch.float32),threshold,True) |
| active=bool(spikes.sum()>0) |
| target=float(goal[0].clamp(.095,.905)) if active else float(state[4]) |
| delta=target-state[4] |
| action=-1 if delta<-.025 else 1 if delta>.025 else 0 |
| return {'action':action,'target_y':target,'decision_active':active,'model':'Target-position LIF SNN', |
| 'spikes':spikes[:,0,:32].int().tolist(),'membrane':mem[:,0,0].tolist(), |
| 'spike_count':int(spikes.sum()),'threshold':threshold,'neurons':64,'steps':10} |
|
|
| def navigate(dx,dy,blocked): |
| with LOCK,torch.inference_mode(): |
| if 'navigation' not in CACHE: |
| net=NavigationPolicy(); net.load_state_dict(torch.load(ROOT/'models'/'navigation.pt',weights_only=True)); net.eval(); CACHE['navigation']=net |
| logits=CACHE['navigation'](torch.tensor([[dx,dy,*blocked]],dtype=torch.float32))[0] |
| |
| for i,b in enumerate(blocked): |
| if b: logits[i]=-1e6 |
| return int(logits.argmax()),logits.softmax(0).tolist() |
|
|
| def transcribe(path): |
| import subprocess,tempfile,soundfile as sf |
| with tempfile.NamedTemporaryFile(suffix='.wav') as out: |
| subprocess.run(['ffmpeg','-v','error','-y','-i',str(path),'-ar','16000','-ac','1',out.name],check=True,capture_output=True) |
| audio,sr=sf.read(out.name) |
| if len(audio)>16000*30: raise ValueError('Please record 30 seconds or less.') |
| with LOCK,torch.inference_mode(): |
| proc,model=load('speech'); inp=proc(audio,sampling_rate=16000,return_tensors='pt').input_features.to(DEVICE) |
| out=model.generate(inp,max_new_tokens=128) |
| return proc.batch_decode(out,skip_special_tokens=True)[0].strip() |
|
|
