# /// script # requires-python = ">=3.11" # dependencies = ["torch==2.8.0", "torchvision==0.23.0", "huggingface-hub>=0.34,<2", "safetensors", "numpy<3", "scipy", "scikit-image", "opencv-python-headless", "pillow", "pyarrow", "requests", "datasets"] # /// import os,sys,time,json,hashlib,shutil,random,math,traceback,io from pathlib import Path import numpy as np import torch from torch import nn import torch.nn.functional as F from torch.utils.data import Dataset,DataLoader from PIL import Image from skimage.color import rgb2lab import cv2,pyarrow as pa,pyarrow.parquet as pq from huggingface_hub import HfApi,hf_hub_download REPO='User-2468/mini-unet-colorizer' BASE='1a9eb8af2754ad2329a24cfe50d388cb559441d0' OLD='experiments/region-coherence-20260928' PREFIX='experiments/recurrent-refinement-20260928' WORK=Path('/work');WORK.mkdir(exist_ok=True);os.chdir(WORK) START=time.monotonic();SEED=92872;STEPS=2000;BATCH=24 api=HfApi(token=os.environ['HF_TOKEN']) assert api.whoami()['name']=='User-2468' HEAD=api.model_info(REPO).sha api.upload_file(repo_id=REPO,path_in_repo=PREFIX+'/access_check.json',path_or_fileobj=json.dumps({'status':'write_access_verified','base_revision':BASE}).encode(),commit_message='Verify experiment artifact write access') print('WRITE_ACCESS_VERIFIED',flush=True) source=WORK/'source';source.mkdir(exist_ok=True) def download(f,rev=HEAD):return Path(hf_hub_download(REPO,f,revision=rev)) sums=json.loads(download(OLD+'/SHA256SUMS.json').read_text()) def verified(f): p=download(OLD+'/'+f) assert hashlib.sha256(p.read_bytes()).hexdigest()==sums[f],f return p for f in ['data.py','legacy_eval.py','metrics.py','model.py','objectives.py','persistence.py','semantic_model.py','spatial.py']: shutil.copy2(verified('source/'+f),source/f) shutil.copy2(verified('data_and_eval.py'),WORK/'r7.py') sys.path.insert(0,str(source));sys.path.insert(0,str(WORK)) RECURRENT_SOURCE="\nimport json\nfrom pathlib import Path\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom safetensors.torch import load_file,save_file\nfrom semantic_model import SemanticColorizer\n\nclass RecurrentColorizer(SemanticColorizer):\n def __init__(self,steps=4,feedback=True,**kwargs):\n super().__init__(**kwargs)\n self.steps=steps;self.feedback=feedback\n self.config.update(architecture='RecurrentColorizer',steps=steps,feedback=feedback,format_version=2)\n def forward(self,L,return_all=False):\n h,w=L.shape[-2:]\n x=F.pad(self.neutral_rgb(L),(0,(-w)%32,0,(-h)%32),mode='replicate')\n x=(x-self.rgb_mean)/self.rgb_std\n features=[]\n for i,layer in enumerate(self.encoder):\n x=layer(x)\n if i in [3,6,12,16]:features.append(x)\n projected=[layer(f) for layer,f in zip(self.lateral,features)]\n x=projected[-1]\n for i in range(2,-1,-1):\n x=self.refine[2-i](F.interpolate(x,size=projected[i].shape[-2:],mode='bilinear',align_corners=False)+projected[i])\n memory=self.memory_norm(torch.cat([F.adaptive_avg_pool2d(f,(8,8)).flatten(2).transpose(1,2) for f in projected[1:]],1))\n q=self.queries[None].expand(L.shape[0],-1,-1)\n px=self.pixel(x);outputs=[];weights=None\n for step in range(self.steps):\n mem=memory\n if step and self.feedback:\n mass=weights.flatten(2).sum(-1,keepdim=True).clamp_min(1e-5)\n pooled=torch.einsum('bqn,bnd->bqd',weights.flatten(2),px.float().flatten(2).transpose(1,2))/mass\n mem=torch.cat([memory,self.memory_norm(pooled.to(memory.dtype))],1)\n old=q\n for block in self.query_blocks:q=block(q,mem)\n if step:q=old+.25*(q-old)\n nq=self.query_norm(q)\n palette=80*torch.tanh(self.palette(nq))\n masks=torch.einsum('bqd,bdhw->bqhw',nq,px)/(self.config['width']**.5)\n weights=F.softmax(masks.float(),dim=1)\n ab=torch.einsum('bqhw,bqc->bchw',weights,palette.float())+2*torch.tanh(self.residual(x).float())\n outputs.append(F.interpolate(ab,size=(x.shape[-2]*4,x.shape[-1]*4),mode='bilinear',align_corners=False)[...,:h,:w])\n return outputs if return_all else outputs[-1]\n\ndef load_recurrent(path,device='cpu'):\n path=Path(path);cfg=json.loads((path/'config.json').read_text())\n m=RecurrentColorizer(**{k:cfg[k] for k in ['head','width','queries','steps','feedback']})\n m.load_state_dict(load_file(str(path/'model.safetensors')))\n return m.to(device).eval()\n" (source/'recurrent_model.py').write_text(RECURRENT_SOURCE) import r7 as r from data import decode_image from semantic_model import load_semantic,save_semantic from recurrent_model import RecurrentColorizer,load_recurrent OUT=r.OUT;r.persistence.PREFIX=PREFIX durable=r.persistence.DurableRun(OUT);r.ev.deadline=lambda:None initial=WORK/'initial';initial.mkdir(exist_ok=True) for f in ['model.safetensors','config.json']:shutil.copy2(download(f,BASE),initial/f) torch.set_num_threads(6);torch.manual_seed(SEED);np.random.seed(SEED);random.seed(SEED) torch.backends.cuda.matmul.allow_tf32=True;torch.backends.cudnn.benchmark=True class TrainPhotos(Dataset): def __init__(self,rgb,targets,regions): self.rgb,self.targets,self.regions=rgb,targets,regions def __len__(self):return len(self.rgb) def __getitem__(self,k): rgb=self.rgb[k].astype(np.float32)/255 if random.random()<.5: g=rgb@np.random.dirichlet([3.,5.,2.]).astype(np.float32) g=np.clip((g**random.uniform(.8,1.25)-.5)*random.uniform(.75,1.15)+.5+random.uniform(-.05,.05),0,1) if random.random()<.4:g=cv2.GaussianBlur(g,(5,5),random.uniform(.3,1.0)) if random.random()<.4:g=np.clip(g+np.random.normal(0,random.uniform(.002,.012),g.shape),0,1) else:g=np.asarray(Image.fromarray(self.rgb[k]).convert('L'),np.float32)/255 L=rgb2lab(np.repeat(g[...,None],3,-1).astype(np.float32))[...,0][None]/50-1 target=self.targets[k].astype(np.float32);regions=self.regions[k].astype(np.int64) if random.random()<.5:L=L[...,::-1];target=target[...,::-1];regions=regions[...,::-1] return torch.from_numpy(L.copy()),torch.from_numpy(target.copy()),torch.from_numpy(regions.copy()) def save(m,path): save_semantic(m,path) if isinstance(m,RecurrentColorizer): (Path(path)/'README.md').write_text('# Experimental recurrent colorizer\n\nRequires source/recurrent_model.py and source/semantic_model.py. Load with load_recurrent(path). Parameters include encoder: 3,994,676. Shared decoder runs four passes by default; encoder runs once. Not production approved.\n') def main(): shutil.copytree(source,OUT/'source',dirs_exist_ok=True,ignore=shutil.ignore_patterns('__pycache__')) shutil.copy2(__file__,OUT/'train.py');shutil.copy2(WORK/'r7.py',OUT/'data_and_eval.py') protocol={'base_revision':BASE,'cache_revision':HEAD,'parameters_including_encoder':3994676,'seed':SEED,'steps_target':STEPS,'batch_size':BATCH, 'arms':['one_pass_control','recurrent_four_pass'], 'architecture':'Reuse pretrained MobileNetV3 and spatial decoder once. Repeat SAME two QueryBlocks; after first pass apply damped residual update .25. Pool spatial features using current color-region weights and append normalized pooled tokens to memory. No extra parameters.', 'loss':'Same fixed selected whole-image targets and losses from region pilot. Control single output; recurrent mean loss across all four passes. No region auxiliary or critic. Matched input batches, initial weights, optimizer and update count.', 'learning_rates':{'encoder':5e-6,'decoder':1e-4},'evaluation':'Fresh 96-image final holdout excluding earlier evaluation IDs and byte-identical training images. Development reused, explicitly nonindependent. 1/2/4 trained-depth and8 extrapolation. Equal-guided-chroma processing. Raw metrics also retained.', 'limitations':['Single seed short architecture pilot','8 passes beyond trained depth may worsen','Fixed targets inherit earlier model selection bias','Does not test diffusion or demonstrate architectural ceiling','No guarantees extra passes help','Near duplicates and upstream pretraining overlap unknown'], 'research':['https://arxiv.org/abs/1806.02919','https://arxiv.org/abs/1904.05290','https://arxiv.org/abs/2212.11613','https://arxiv.org/abs/2111.05826'], 'budget':{'hardware':'l4x1','hard_timeout_minutes':30,'maximum_hardware_usd_at_published_rate':.40,'training_graceful_stop_minutes':21}, 'release':'Experimental checkpoints to main subfolder, visual review required before production'} r.write('protocol.json',protocol) durable.sync('protocol and sources before GPU training') # Validate equivalent first-pass computation on CPU before enabling GPU kernels. baseline=load_semantic(initial,'cpu').requires_grad_(False) control=load_semantic(initial,'cpu') recurrent=RecurrentColorizer(steps=4).cpu() recurrent.load_state_dict(control.state_dict(),strict=True) assert sum(p.numel() for p in recurrent.parameters())==3994676 recurrent.eval() checks={} with torch.no_grad(): for h,w in [(96,128),(97,113)]: x=torch.rand(2,1,h,w)*2-1 ys=recurrent(x,return_all=True);bp=baseline(x) checks[f'cpu_{h}x{w}_max_abs']=float((ys[0]-bp).abs().max()) assert torch.allclose(ys[0],bp,atol=2e-5),checks assert all(torch.isfinite(y).all() for y in ys) baseline=baseline.cuda();control=control.cuda();recurrent=recurrent.cuda() x=torch.rand(2,1,96,128,device='cuda')*2-1 # Strict parity checks require full precision, not approximate TF32 kernels. torch.backends.cuda.matmul.allow_tf32=False torch.backends.cudnn.allow_tf32=False torch.backends.cudnn.benchmark=False torch.backends.cudnn.deterministic=True with torch.no_grad(): ys=recurrent(x,return_all=True);bp=baseline(x) checks['gpu_fp32_max_abs']=float((ys[0]-bp).abs().max()) print('PARITY_CHECKS',json.dumps(checks),flush=True) assert torch.allclose(ys[0],bp,atol=2e-5),checks recurrent(x).square().mean().backward() assert all(torch.isfinite(p.grad).all() for p in recurrent.parameters() if p.grad is not None) recurrent.zero_grad(set_to_none=True) save(recurrent,OUT/'initial_recurrent') test=load_recurrent(OUT/'initial_recurrent','cuda') with torch.no_grad():assert torch.allclose(test(x),recurrent(x),atol=1e-5) del test r.write('smoke.json',{'passed':True,'first_pass_matches_v3':True,'parity':checks,'finite_recurrent_backward':True,'reload_equal':True,'parameter_count':3994676}) durable.sync('CPU and FP32 GPU parity, backward and reload verified') print('PREFLIGHT_PASSED',flush=True) torch.backends.cuda.matmul.allow_tf32=True torch.backends.cudnn.allow_tf32=True torch.backends.cudnn.benchmark=True torch.backends.cudnn.deterministic=False tables=[] for repo,rev,files in [('johnowhitaker/imagenette2-320','771c1310a2487e8076ede6b7d6307244aa8400af',['default/train/0000.parquet']),('detection-datasets/coco','26ddc382fe75dfc2a0655b5977e296ea10efebce',['default/train/0000.parquet','default/train/0001.parquet'])]: for f in files:tables.append(pq.read_table(hf_hub_download(repo,f,repo_type='dataset',revision=rev),columns=['image'])) table=pa.concat_tables(tables) manifest=json.loads(verified('manifest.json').read_text());ids=manifest['train'] targets=np.empty((len(ids),2,64,64),np.float16);regions=np.empty((len(ids),64,64),np.uint8) for start in range(0,len(ids),1024): with np.load(verified(f'region_cache/{start:05d}.npz')) as z: assert z['indices'].tolist()==ids[start:start+len(z['indices'])] targets[start:start+len(z['indices'])]=z['target'];regions[start:start+len(z['indices'])]=z['regions'] rgbcache=np.stack([np.asarray(decode_image(table,i).convert('RGB').resize((256,256),Image.Resampling.BILINEAR)) for i in ids]) fresh=pq.read_table(hf_hub_download('detection-datasets/coco','default/val/0001.parquet',repo_type='dataset',revision=manifest['coco_revision']),columns=['image']) excludes=set(manifest['development_regression']+manifest['heldout_regression']+manifest['new_holdout']) oldeval=json.loads(download('experiments/round7-20260928/manifest.json',BASE).read_text()) excludes.update(oldeval['development']+oldeval['heldout']) oldtrain=json.loads(download('experiments/region-coherence-20260928/source/previous_manifest.json').read_text())['train'] hashes={hashlib.sha256(table['image'][i].as_py()['bytes']).digest() for i in oldtrain} hold=[int(i) for i in np.random.default_rng(SEED).permutation(len(fresh)) if int(i) not in excludes and hashlib.sha256(fresh['image'][int(i)].as_py()['bytes']).digest() not in hashes][:96] assert len(hold)==96 dev=manifest['development_regression'][:32] r.write('manifest.json',{'train':ids,'development_regression':dev,'fresh_holdout':hold,'source_manifest_revision':HEAD,'seed':SEED}) models={'control':control,'recurrent':recurrent} opts={n:torch.optim.AdamW([{'params':m.encoder.parameters(),'lr':5e-6,'base_lr':5e-6},{'params':[p for k,p in m.named_parameters() if not k.startswith('encoder.')],'lr':1e-4,'base_lr':1e-4}],weight_decay=.01) for n,m in models.items()} loader=DataLoader(TrainPhotos(rgbcache,targets,regions),batch_size=BATCH,shuffle=True,num_workers=4,pin_memory=True,drop_last=True,persistent_workers=True,generator=torch.Generator().manual_seed(SEED)) it=iter(loader);history=[];last=0 def persist(step): for n,m in models.items(): save(m,OUT/n/f'step_{step:05d}') torch.save({'step':step,'optimizer':opts[n].state_dict(),'torch_rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all(),'numpy_rng':np.random.get_state(),'python_rng':random.getstate(),'exact_dataloader_resume':False},OUT/n/'optimizer.pt') r.write('history.json',history);r.write('status.json',{'status':'training','step':step,'elapsed_seconds':time.monotonic()-START,'production_approved':False}) durable.sync('matched recurrent trial step '+str(step)) durable.sync('cached targets verified and holdout frozen') for step in range(1,STEPS+1): if time.monotonic()-START>21*60:break try:L,t,_=next(it) except StopIteration:it=iter(loader);L,t,_=next(it) L,t=L.cuda(non_blocking=True),t.cuda(non_blocking=True);row={'step':step} for n,m in models.items(): m.train() for layer in m.encoder.modules(): if isinstance(layer,nn.BatchNorm2d):layer.eval() opts[n].zero_grad(set_to_none=True) factor=min(step/100,1)*(.2+.8*(1+math.cos(math.pi*step/STEPS))/2) for group in opts[n].param_groups:group['lr']=group['base_lr']*factor with torch.autocast('cuda',dtype=torch.bfloat16): ps=m(L,return_all=True) if n=='recurrent' else [m(L)] terms=torch.stack([r.losses(p,t).mean() for p in ps]);loss=terms.mean() assert torch.isfinite(loss) loss.backward();torch.nn.utils.clip_grad_norm_(m.parameters(),1,error_if_nonfinite=True);opts[n].step() row[n]=[float(v.detach()) for v in terms] last=step if step%100==0: row['elapsed_seconds']=time.monotonic()-START;history.append(row);print('STEP',json.dumps(row),flush=True) if step%500==0:persist(step) assert last>0,'No training budget remained' persist(last) for n,m in models.items():save(m,OUT/n/'candidate') class View(nn.Module): def __init__(self,m,steps):super().__init__();self.m=m;self.steps=steps def forward(self,x): old=self.m.steps;self.m.steps=self.steps try:return self.m(x) finally:self.m.steps=old def decode(self,x,temperature=.38):return x views={'v3':baseline,'control':control} views.update({f'recurrent_{s}':View(recurrent,s) for s in [1,2,4,8]}) for n,m in views.items(): m.eval() r.write(n+'_fresh.json',r.ev.evaluate(m,fresh,hold)) r.ev.probes(m,n+'_probes') x=torch.zeros(1,1,256,256,device='cuda') with torch.inference_mode(): for _ in range(3):m(x) torch.cuda.synchronize();a=time.monotonic() for _ in range(20):m(x) torch.cuda.synchronize() r.write(n+'_latency.json',{'batch':1,'size':256,'precision':'float32','ms_per_image':(time.monotonic()-a)*1000/20,'hardware':torch.cuda.get_device_name()}) durable.sync(n+' final metrics and probes') r.grid({'v3':baseline,'control':control,'recurrent_4':views['recurrent_4']},fresh,hold,'fresh_architecture') r.grid({n:views[n] for n in ['recurrent_1','recurrent_2','recurrent_4','recurrent_8']},fresh,hold,'fresh_passes') r.grid({'v3':baseline,'control':control,'recurrent_4':views['recurrent_4']},fresh,dev,'development') import requests for name,url in {'migrant_mother':'https://cdn.loc.gov/service/pnp/ppmsca/50200/50236v.jpg','power_house_mechanic':'https://www.archives.gov/exhibits/picturing_the_century/images/port_hine_022_v36.jpg'}.items(): try: response=requests.get(url,timeout=30);response.raise_for_status() tb=pa.table({'image':[{'bytes':response.content,'path':None}]}) r.grid({'v3':baseline,'control':control,'recurrent_4':views['recurrent_4']},tb,[0],name) r.write(name+'_source.json',{'url':url,'sha256':hashlib.sha256(response.content).hexdigest(),'previously_seen_probe':True}) except Exception as e:r.write(name+'_error.json',{'type':type(e).__name__}) r.write('status.json',{'status':'completed','steps':last,'elapsed_seconds':time.monotonic()-START,'production_approved':False,'next':'Inspect images and 1/2/4/8 pass progression. Eight passes are extrapolation, not promised improvement.'}) durable.sync('recurrent experiment completed with visual comparisons') if __name__=='__main__': try:main() except BaseException: msg=traceback.format_exc().replace(os.environ.get('HF_TOKEN','__NONE__'),'[REDACTED]') r.write('failure.json',{'traceback':msg}) try:durable.sync('failure report') except Exception:pass print(msg,flush=True) raise SystemExit(1)