| """Single / multiple coherent palettes, optional compact global-attention encoder. |
| |
| Original experiment inspired by diverse colorization and multiple-choice learning; |
| not a reproduction. All deployed parameters, including the selector, are counted. |
| """ |
| import json |
| from pathlib import Path |
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
| from safetensors.torch import save_file, load_file |
| from semantic_model import SemanticColorizer |
|
|
|
|
| class DecisionColorizer(SemanticColorizer): |
| def __init__(self, modes=1, backbone='mobilenet', pretrained=False, |
| encoder_revision=None, **kwargs): |
| super().__init__(pretrained=False, **kwargs) |
| self.modes = modes |
| self.backbone = backbone |
| d = self.config['width']; q = self.config['queries'] |
| if backbone == 'mobilevit': |
| import timm |
| from huggingface_hub import hf_hub_download |
| name = 'mobilevitv2_075.cvnets_in1k' |
| overlay = None |
| if pretrained: |
| assert encoder_revision, 'Pin pretrained encoder revision' |
| filename = hf_hub_download('timm/'+name, 'model.safetensors', revision=encoder_revision) |
| overlay = {'file': filename, 'hf_hub_id': None} |
| self.encoder = timm.create_model(name, pretrained=pretrained, |
| pretrained_cfg_overlay=overlay, features_only=True, out_indices=(1,2,3,4)) |
| self.lateral = nn.ModuleList([nn.Conv2d(c,d,1) for c in self.encoder.feature_info.channels()]) |
| cfg = self.encoder.pretrained_cfg |
| self.rgb_mean.copy_(torch.tensor(cfg['mean']).view(1,3,1,1)) |
| self.rgb_std.copy_(torch.tensor(cfg['std']).view(1,3,1,1)) |
| if modes > 1: |
| self.mode_embeddings = nn.Parameter(torch.randn(modes,d)*.03) |
| self.mode_offsets = nn.Parameter(torch.zeros(modes,q,2)) |
| |
| with torch.no_grad(): |
| a=torch.arange(modes)*2*torch.pi/modes |
| self.mode_offsets[:,:,0]=.025*a.cos()[:,None] |
| self.mode_offsets[:,:,1]=.025*a.sin()[:,None] |
| self.mode_score = nn.Linear(d,modes) |
| nn.init.zeros_(self.mode_score.weight);nn.init.zeros_(self.mode_score.bias) |
| self.config.update(architecture='DecisionColorizer', format_version=3, |
| modes=modes, backbone=backbone, encoder_revision=encoder_revision) |
| count = sum(p.numel() for p in self.parameters()) |
| assert count < 4_000_000, count |
|
|
| def features(self,L): |
| h,w=L.shape[-2:] |
| x=F.pad(self.neutral_rgb(L),(0,(-w)%32,0,(-h)%32),mode='replicate') |
| x=(x-self.rgb_mean)/self.rgb_std |
| if self.backbone=='mobilevit': |
| fs=self.encoder(x) |
| else: |
| fs=[] |
| for i,layer in enumerate(self.encoder): |
| x=layer(x) |
| if i in [3,6,12,16]:fs.append(x) |
| projected=[layer(f) for layer,f in zip(self.lateral,fs)] |
| x=projected[-1] |
| for i in range(2,-1,-1): |
| x=self.refine[2-i](F.interpolate(x,size=projected[i].shape[-2:],mode='bilinear',align_corners=False)+projected[i]) |
| memory=self.memory_norm(torch.cat([F.adaptive_avg_pool2d(f,(8,8)).flatten(2).transpose(1,2) for f in projected[1:]],1)) |
| q=self.queries[None].expand(L.shape[0],-1,-1) |
| for block in self.query_blocks:q=block(q,memory) |
| q=self.query_norm(q) |
| masks=torch.einsum('bqd,bdhw->bqhw',q,self.pixel(x))/(self.config['width']**.5) |
| weights=F.softmax(masks.float(),dim=1) |
| residual=2*torch.tanh(self.residual(x).float()) |
| return q,weights,residual,memory.mean(1) |
|
|
| def forward(self,L,return_details=False): |
| q,weights,residual,context=self.features(L) |
| if self.modes==1: |
| palette=80*torch.tanh(self.palette(q))[:,None] |
| scores=q.new_zeros((len(L),1)) |
| else: |
| mq=q[:,None]+self.mode_embeddings[None,:,None] |
| palette=80*torch.tanh(self.palette(mq)+self.mode_offsets[None]) |
| scores=self.mode_score(context) |
| outputs=torch.einsum('bqhw,bkqc->bkchw',weights,palette.float())+residual[:,None] |
| b,k,c,h,w=outputs.shape |
| full=F.interpolate(outputs.reshape(b*k,c,h,w),size=(h*4,w*4),mode='bilinear',align_corners=False) |
| full=full[...,:L.shape[-2],:L.shape[-1]].reshape(b,k,c,*L.shape[-2:]) |
| default=full[torch.arange(b,device=L.device),scores.argmax(1)] |
| if return_details: |
| return {'output':default,'all':full,'scores':scores.float(), |
| 'weights':weights,'palette':palette.float()} |
| return default |
|
|
|
|
| def initialize_from_v3(model,baseline): |
| old=baseline.state_dict();state=model.state_dict() |
| copied=[] |
| for k,v in old.items(): |
| if model.backbone=='mobilevit' and (k.startswith(('encoder.','lateral.')) or k in ['rgb_mean','rgb_std']):continue |
| if k in state and v.shape==state[k].shape: |
| state[k]=v;copied.append(k) |
| model.load_state_dict(state,strict=True) |
| return copied |
|
|
|
|
| def save_decision(m,path): |
| p=Path(path);p.mkdir(parents=True,exist_ok=True) |
| save_file({k:v.detach().cpu().contiguous() for k,v in m.state_dict().items()},str(p/'model.safetensors')) |
| (p/'config.json').write_text(json.dumps(m.config,indent=2)) |
|
|
|
|
| def load_decision(path,device='cpu'): |
| p=Path(path);cfg=json.loads((p/'config.json').read_text()) |
| m=DecisionColorizer(**{k:cfg[k] for k in ['modes','backbone','width','queries','head','encoder_revision']}) |
| m.load_state_dict(load_file(str(p/'model.safetensors')),strict=True) |
| return m.to(device).eval() |
|
|
