File size: 4,943 Bytes
704fa80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """Shared, checkpoint-compatible Mini U-Net. Parameters remain under 4M."""
import json
import math
from pathlib import Path
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from huggingface_hub import PyTorchModelHubMixin, snapshot_download
from safetensors.torch import load_file, save_file
def double_conv(in_ch, out_ch):
return nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
)
class DilatedContextBlock(nn.Module):
def __init__(self, channels, mid_ch=96, dilations=(2, 4, 8)):
super().__init__()
self.proj_in = nn.Sequential(
nn.Conv2d(channels, mid_ch, 1, bias=False),
nn.BatchNorm2d(mid_ch),
nn.ReLU(inplace=True),
)
layers = []
for d in dilations:
layers += [
nn.Conv2d(mid_ch, mid_ch, 3, padding=d, dilation=d, bias=False),
nn.BatchNorm2d(mid_ch),
nn.ReLU(inplace=True),
]
self.dilated = nn.Sequential(*layers)
self.proj_out = nn.Sequential(
nn.Conv2d(mid_ch, channels, 1, bias=False),
nn.BatchNorm2d(channels),
)
nn.init.zeros_(self.proj_out[-1].weight)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
y = self.proj_in(x)
y = self.dilated(y)
y = self.proj_out(y)
return self.relu(x + y)
class SmallUNetColorizer(
nn.Module,
PyTorchModelHubMixin,
pipeline_tag="image-to-image",
license="apache-2.0",
tags=["colorization", "unet", "image-to-image", "classification"],
):
def __init__(self, bin_centers, in_ch: int = 1, base: int = 44,
context_mid_ch: int = 96, context_dilations=(2, 4, 8)):
super().__init__()
self.in_ch, self.base = in_ch, base
num_bins = len(bin_centers)
self.num_bins = num_bins
self.register_buffer("bin_centers", torch.tensor(bin_centers, dtype=torch.float32))
self.enc1 = double_conv(in_ch, base)
self.enc2 = double_conv(base, base * 2)
self.enc3 = double_conv(base * 2, base * 4)
self.enc4 = double_conv(base * 4, base * 8)
self.pool = nn.MaxPool2d(2)
self.context = DilatedContextBlock(base * 8, mid_ch=context_mid_ch,
dilations=tuple(context_dilations))
self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2)
self.dec3 = double_conv(base * 8, base * 4)
self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2)
self.dec2 = double_conv(base * 4, base * 2)
self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2)
self.dec1 = double_conv(base * 2, base)
self.out_conv = nn.Conv2d(base, num_bins, 1)
def forward(self, x):
h, w = x.shape[-2:]
x = F.pad(x, (0, (-w) % 8, 0, (-h) % 8), mode="replicate")
e1 = self.enc1(x)
e2 = self.enc2(self.pool(e1))
e3 = self.enc3(self.pool(e2))
e4 = self.context(self.enc4(self.pool(e3)))
d3 = self.dec3(torch.cat([self.up3(e4), e3], dim=1))
d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
return self.out_conv(d1)[..., :h, :w]
def decode(self, logits, temperature: float = 0.38):
if not math.isfinite(temperature) or temperature <= 0:
raise ValueError("temperature must be finite and positive")
probs_t = F.softmax(logits.float() / temperature, dim=1)
return torch.einsum("bqhw,qc->bchw", probs_t, self.bin_centers)
def load_model(source, revision=None, device="cpu"):
path = Path(source)
if not path.is_dir():
path = Path(snapshot_download(source, revision=revision,
allow_patterns=["config.json", "model.safetensors"]))
config = json.loads((path / "config.json").read_text())
state = load_file(str(path / "model.safetensors"))
centers = torch.tensor(config["bin_centers"], dtype=torch.float32)
if not torch.equal(centers, state["bin_centers"]):
raise ValueError("Checkpoint config and state color bins differ; refusing ambiguous decode")
model = SmallUNetColorizer(**config)
model.load_state_dict(state, strict=True)
model.to(device).eval()
return model
def save_model(model, path):
path = Path(path); path.mkdir(parents=True, exist_ok=True)
model.save_pretrained(path)
# Mixin config can retain constructor bins; use the actual authoritative buffer.
cfg = json.loads((path / "config.json").read_text())
cfg["bin_centers"] = model.bin_centers.detach().cpu().tolist()
(path / "config.json").write_text(json.dumps(cfg, indent=2))
|