| """Compact pretrained semantic colorizer; dense and shared-palette variants.""" |
| import json |
| from pathlib import Path |
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
| from torchvision.models import mobilenet_v3_large, MobileNet_V3_Large_Weights |
| from safetensors.torch import save_file, load_file |
|
|
| class QueryBlock(nn.Module): |
| def __init__(self,d=96): |
| super().__init__() |
| self.self_attn=nn.MultiheadAttention(d,4,batch_first=True,dropout=0) |
| self.cross_attn=nn.MultiheadAttention(d,4,batch_first=True,dropout=0) |
| self.norms=nn.ModuleList([nn.LayerNorm(d) for _ in range(3)]) |
| self.ff=nn.Sequential(nn.Linear(d,2*d),nn.GELU(),nn.Linear(2*d,d)) |
| def forward(self,q,memory): |
| x=self.norms[0](q);q=q+self.self_attn(x,x,x,need_weights=False)[0] |
| x=self.norms[1](q);q=q+self.cross_attn(x,memory,memory,need_weights=False)[0] |
| return q+self.ff(self.norms[2](q)) |
|
|
| def refine(d): |
| return nn.Sequential(nn.Conv2d(d,d,3,padding=1,bias=False),nn.GroupNorm(8,d),nn.SiLU()) |
|
|
| class SemanticColorizer(nn.Module): |
| def __init__(self,head='palette',pretrained=False,width=128,queries=16): |
| super().__init__() |
| if head not in ['palette','dense']:raise ValueError(head) |
| self.config={'architecture':'SemanticColorizer','head':head,'width':width,'queries':queries,'format_version':1} |
| self.encoder=mobilenet_v3_large(weights=MobileNet_V3_Large_Weights.IMAGENET1K_V2 if pretrained else None,progress=False).features |
| self.lateral=nn.ModuleList([nn.Conv2d(c,width,1) for c in [24,40,112,960]]) |
| self.refine=nn.ModuleList([refine(width) for _ in range(3)]) |
| self.register_buffer('rgb_mean',torch.tensor([.485,.456,.406]).view(1,3,1,1)) |
| self.register_buffer('rgb_std',torch.tensor([.229,.224,.225]).view(1,3,1,1)) |
| if head=='palette': |
| self.queries=nn.Parameter(torch.randn(queries,width)*.2) |
| self.query_blocks=nn.ModuleList([QueryBlock(width) for _ in range(2)]) |
| self.memory_norm=nn.LayerNorm(width) |
| self.query_norm=nn.LayerNorm(width) |
| self.pixel=nn.Conv2d(width,width,1) |
| self.palette=nn.Sequential(nn.Linear(width,width),nn.GELU(),nn.Linear(width,2)) |
| self.residual=nn.Conv2d(width,2,1) |
| nn.init.normal_(self.palette[-1].weight,std=.01);nn.init.zeros_(self.palette[-1].bias) |
| nn.init.zeros_(self.residual.weight);nn.init.zeros_(self.residual.bias) |
| else: |
| self.dense=nn.Sequential(refine(width),nn.Conv2d(width,2,1)) |
| nn.init.normal_(self.dense[-1].weight,std=.01);nn.init.zeros_(self.dense[-1].bias) |
| count=sum(p.numel() for p in self.parameters()) |
| if count>=4_000_000:raise ValueError(f'Parameter budget exceeded: {count}') |
|
|
| @staticmethod |
| def neutral_rgb(L): |
| light=(L.float()*50+50).clamp(0,100) |
| y=torch.where(light>8,((light+16)/116)**3,light/903.296296) |
| g=torch.where(y<=.0031308,12.92*y,1.055*y.clamp_min(1e-8).pow(1/2.4)-.055) |
| return g.expand(-1,3,-1,-1) |
|
|
| def forward(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 |
| features=[] |
| for i,layer in enumerate(self.encoder): |
| x=layer(x) |
| if i in [3,6,12,16]:features.append(x) |
| projected=[layer(f) for layer,f in zip(self.lateral,features)] |
| 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]) |
| if self.config['head']=='palette': |
| memory=torch.cat([F.adaptive_avg_pool2d(f,(8,8)).flatten(2).transpose(1,2) for f in projected[1:]],1) |
| memory=self.memory_norm(memory) |
| 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) |
| palette=80*torch.tanh(self.palette(q)) |
| masks=torch.einsum('bqd,bdhw->bqhw',q,self.pixel(x))/(self.config['width']**.5) |
| weights=F.softmax(masks.float(),dim=1) |
| ab=torch.einsum('bqhw,bqc->bchw',weights,palette.float())+2*torch.tanh(self.residual(x).float()) |
| else:ab=80*torch.tanh(self.dense(x).float()) |
| return F.interpolate(ab,size=(x.shape[-2]*4,x.shape[-1]*4),mode='bilinear',align_corners=False)[...,:h,:w] |
| def decode(self,z,temperature=.38):return z.float() |
|
|
| def save_semantic(model,path): |
| path=Path(path);path.mkdir(parents=True,exist_ok=True) |
| save_file({k:v.detach().cpu().contiguous() for k,v in model.state_dict().items()},str(path/'model.safetensors')) |
| (path/'config.json').write_text(json.dumps(model.config,indent=2)) |
| (path/'README.md').write_text('# Experimental semantic colorizer\n\nHead: '+model.config['head']+'. Parameters: '+str(sum(p.numel() for p in model.parameters()))+'.\n\nNot approved for production. Requires semantic_model.py; incompatible with the old U-Net loader. Input is Lab lightness normalized to [-1,1]; output is Lab ab. See the run protocol, provenance, selection and visual comparisons. Predictions are plausible colors, not recovered historical truth.\n') |
|
|
| def load_semantic(path,device='cpu'): |
| path=Path(path);cfg=json.loads((path/'config.json').read_text()) |
| model=SemanticColorizer(**{k:cfg[k] for k in ['head','width','queries']}) |
| model.load_state_dict(load_file(str(path/'model.safetensors')),strict=True) |
| return model.to(device).eval() |
|
|
