| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """ |
| Train a tiny (<4M parameter) U-Net that colorizes grayscale photos. |
| |
| Classification-style (Zhang et al., "Colorful Image Colorization"): predicts |
| a distribution over quantized Lab a/b color bins per pixel, rather than |
| regressing a single ab value directly. Bin classes are up-weighted by rarity |
| in the real training-data color distribution, so the model isn't rewarded |
| for hedging toward desaturated averages the way a plain L1/L2 regression |
| loss would. Decode with an annealed mean (see `SmallUNetColorizer.decode`). |
| |
| Designed to run as a Hugging Face Job: |
| |
| hf jobs uv run --flavor t4-small --timeout 1h \ |
| -s HF_TOKEN colorize_train.py \ |
| --push-to-hub --hub-model-id <your-username>/mini-unet-colorizer |
| |
| Run locally the same way, just with `uv run colorize_train.py ...` or |
| `python colorize_train.py ...` after `pip install`-ing the deps above. |
| """ |
| import argparse |
| import os |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from huggingface_hub import HfApi, PyTorchModelHubMixin, hf_hub_download |
| from safetensors.torch import load_file as load_safetensors |
| from scipy.ndimage import gaussian_filter |
| from scipy.spatial import cKDTree |
| from skimage.color import lab2rgb, rgb2lab |
| from torch.utils.data import DataLoader, Dataset |
| from torchvision import transforms as T |
| from torchvision.utils import make_grid, save_image |
| from tqdm.auto import tqdm |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| class ColorBins: |
| """The quantized ab bin grid, its empirical prior, and rebalancing weights. |
| |
| `bin_centers`: (Q, 2) array of [a, b] centers for the Q bins that actually |
| occur in the sampled data (a data-driven stand-in for "in-gamut" bins, |
| rather than computing exact sRGB gamut geometry). |
| `weights`: (Q,) rebalancing weights, high for rare/saturated bins, low |
| for common/desaturated ones, E_p[weight] == 1. |
| """ |
|
|
| def __init__(self, bin_centers: np.ndarray, weights: np.ndarray): |
| assert bin_centers.shape[0] == weights.shape[0] |
| self.bin_centers = bin_centers.astype(np.float32) |
| self.weights = weights.astype(np.float32) |
| self.num_bins = bin_centers.shape[0] |
| self._tree = cKDTree(self.bin_centers) |
|
|
| def soft_encode(self, ab: np.ndarray, k: int = 5, sigma: float = 5.0): |
| """ab: (H, W, 2) continuous Lab a/b values. |
| Returns (top_idx, top_weight): each (H, W, k). top_weight sums to 1 |
| over the last axis (a Gaussian-kernel soft target over the k nearest |
| bins, per Zhang et al.'s soft-encoding scheme).""" |
| h, w, _ = ab.shape |
| flat = ab.reshape(-1, 2) |
| dist, idx = self._tree.query(flat, k=k) |
| if k == 1: |
| dist, idx = dist[:, None], idx[:, None] |
| weight = np.exp(-(dist ** 2) / (2 * sigma ** 2)) |
| weight /= weight.sum(axis=1, keepdims=True) |
| return idx.reshape(h, w, k).astype(np.int64), weight.reshape(h, w, k).astype(np.float32) |
|
|
|
|
| def image_mean_chroma(lab): |
| """Mean Lab chroma (sqrt(a^2+b^2)) -- near 0 for a genuinely grayscale- |
| source image (real color photos, even desaturated/overcast ones, have |
| meaningfully more residual chroma than a true B&W photo's compression |
| noise).""" |
| return float(np.sqrt(lab[:, :, 1] ** 2 + lab[:, :, 2] ** 2).mean()) |
|
|
|
|
| def build_color_bins(hf_split, num_samples=3000, grid_size=10, ab_range=110, |
| rebalance_lambda=0.5, smoothing_sigma=1.0, seed=0, |
| grayscale_chroma_threshold=3.0): |
| """Derive the bin set + rebalancing weights from a sample of real training |
| images, by actually converting them to Lab and histogramming ab values on |
| a `grid_size`-spaced grid. Bins that never occur simply aren't included -- |
| this is a cheap data-driven stand-in for restricting to in-gamut bins. |
| |
| Images whose mean chroma is below `grayscale_chroma_threshold` are |
| excluded from the histogram: a genuinely black-and-white source photo's |
| correct ab really is ~0, which is a different phenomenon from the model |
| hedging toward gray on a photo that actually had color. Mixing the two |
| into one "near-zero is common" statistic would miscalibrate the |
| rebalancing weights -- overcounting how often gray is truly correct. |
| """ |
| edges = np.arange(-ab_range, ab_range + grid_size, grid_size) |
| centers_1d = (edges[:-1] + edges[1:]) / 2 |
| n_grid = len(centers_1d) |
|
|
| rng = np.random.default_rng(seed) |
| n = len(hf_split) |
| sample_idx = rng.choice(n, size=min(num_samples, n), replace=False) |
|
|
| counts = np.zeros((n_grid, n_grid), dtype=np.int64) |
| n_grayscale_source = 0 |
| for i in sample_idx: |
| img = hf_split[int(i)]["image"] |
| if img.mode != "RGB": |
| img = img.convert("RGB") |
| img = img.resize((64, 64)) |
| arr = np.asarray(img).astype(np.float32) / 255.0 |
| lab = rgb2lab(arr) |
| if image_mean_chroma(lab) < grayscale_chroma_threshold: |
| n_grayscale_source += 1 |
| continue |
| a_idx = np.clip(((lab[:, :, 1] + ab_range) // grid_size).astype(int), 0, n_grid - 1) |
| b_idx = np.clip(((lab[:, :, 2] + ab_range) // grid_size).astype(int), 0, n_grid - 1) |
| np.add.at(counts, (a_idx.ravel(), b_idx.ravel()), 1) |
|
|
| frac = n_grayscale_source / len(sample_idx) |
| print(f"{n_grayscale_source}/{len(sample_idx)} ({frac:.1%}) sampled images are " |
| f"grayscale-source (mean chroma < {grayscale_chroma_threshold}) -- excluded from " |
| f"the color-bin statistics below.") |
|
|
| smoothed = gaussian_filter(counts.astype(np.float32), sigma=smoothing_sigma) |
|
|
| keep_mask = counts > 0 |
| keep_a, keep_b = np.nonzero(keep_mask) |
| bin_centers = np.stack([centers_1d[keep_a], centers_1d[keep_b]], axis=1) |
|
|
| prior = smoothed[keep_a, keep_b] |
| prior = prior / prior.sum() |
| q = len(prior) |
| w = 1.0 / ((1 - rebalance_lambda) * prior + rebalance_lambda / q) |
| w = w / (prior * w).sum() |
|
|
| return ColorBins(bin_centers, w) |
|
|
|
|
| |
| |
| |
| |
| 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): |
| """Cheap receptive-field expansion at the bottleneck. The plain U-Net's |
| receptive field is fixed in absolute pixels regardless of input |
| resolution (~51-54px, empirically measured), so at higher --image-size |
| it covers a shrinking fraction of the image -- the likely cause of |
| color decisions that don't stay consistent across a whole object (e.g. |
| a parachute canopy getting inconsistent patches of color). |
| |
| A naive fix (another U-Net level, channel-doubling as the existing ones |
| do) costs ~4.5M params just for one extra block -- unaffordable. This |
| instead projects down to a narrow `mid_ch` width, runs a few dilated 3x3 |
| convs (same param cost as a normal 3x3 conv, much larger spatial reach: |
| empirically ~232x232px in isolation with dilations (2,4,8), verified via |
| gradient measurement), and projects back up. Residual connection, so it |
| starts as a small perturbation on top of the existing backbone (which |
| matters for warm-starting from a checkpoint that predates this block -- |
| it won't disrupt what's already learned, and its influence grows through |
| training as it learns to use the wider context where it helps) rather |
| than a disruptive change. |
| """ |
|
|
| 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), |
| ) |
| 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"], |
| ): |
| """U-Net colorizer, classification-style (Zhang et al.). In: L in [-1,1], |
| shape (B,1,H,W). Out: per-pixel logits over `len(bin_centers)` quantized |
| ab bins, shape (B,Q,H,W) -- decode with `.decode()` to get back ab. |
| |
| `bin_centers` (a plain list of [a, b] pairs, JSON-serializable so it |
| round-trips through config.json) is stored as a buffer, so every |
| checkpoint carries its own bin grid -- no separate file to keep in sync. |
| """ |
|
|
| 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): |
| 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) |
|
|
| def decode(self, logits, temperature: float = 0.38): |
| """Annealed-mean decode: T=1 is the full expectation (smooth but |
| desaturated -- the same hedging problem as regression); T->0 |
| approaches the argmax/mode (vivid but can be noisy/blotchy per |
| pixel). T=0.38 (Zhang et al.'s value) balances the two. Returns raw |
| ab in Lab units, shape (B, 2, H, W).""" |
| logp = F.log_softmax(logits, dim=1) |
| probs_t = F.softmax(logp / temperature, dim=1) |
| return torch.einsum("bqhw,qc->bchw", probs_t, self.bin_centers) |
|
|
|
|
| |
| |
| |
| class ColorizationDataset(Dataset): |
| """Wraps a HF `datasets` split that has an 'image' column of PIL images. |
| Returns L (continuous, normalized), continuous ab (for visualization |
| ground truth only), and the soft-encoded bin targets used for the loss.""" |
|
|
| def __init__(self, hf_split, color_bins, image_size=128, augment=True): |
| self.ds = hf_split |
| self.color_bins = color_bins |
| aug = [T.RandomHorizontalFlip()] if augment else [] |
| self.resize = T.Compose([T.Resize((image_size, image_size)), *aug]) |
|
|
| def __len__(self): |
| return len(self.ds) |
|
|
| def __getitem__(self, idx): |
| img = self.ds[idx]["image"] |
| if img.mode != "RGB": |
| img = img.convert("RGB") |
| img = self.resize(img) |
| arr = np.asarray(img).astype(np.float32) / 255.0 |
| lab = rgb2lab(arr).astype(np.float32) |
| L = lab[:, :, 0:1] / 50.0 - 1.0 |
| ab = lab[:, :, 1:3] |
|
|
| top_idx, top_weight = self.color_bins.soft_encode(ab) |
|
|
| L = torch.from_numpy(np.ascontiguousarray(L.transpose(2, 0, 1))) |
| ab = torch.from_numpy(np.ascontiguousarray(ab.transpose(2, 0, 1))) |
| top_idx = torch.from_numpy(top_idx) |
| top_weight = torch.from_numpy(top_weight) |
| return L, ab, top_idx, top_weight |
|
|
|
|
| def classification_loss(logits, top_idx, top_weight, class_weights): |
| """Weighted soft multinomial cross-entropy (Zhang et al.). |
| |
| logits: (B,Q,H,W) raw model output. |
| top_idx, top_weight: (B,H,W,5) from ColorBins.soft_encode -- the soft |
| target distribution over each pixel's 5 nearest bins. |
| class_weights: (Q,) rebalancing weights (high for rare/saturated bins), |
| applied per-pixel using that pixel's single nearest bin. |
| |
| This is the classification analogue of the earlier `--color-weight` |
| trick, but built into the loss's class-rebalancing rather than bolted on |
| after the fact: rare, saturated bins get more weight *because* they're |
| rare in the data, computed directly from the empirical color |
| distribution instead of a hand-picked multiplier. |
| |
| Deliberately avoids ever materializing a full (B,Q,H,W) log-softmax |
| tensor -- we only need log-probabilities at 5 of the Q bins per pixel, |
| not all of them, and with Q in the hundreds that full tensor is the |
| single biggest activation in the whole model (bigger than anything in |
| the backbone), so this matters a lot for fitting on a 16GB GPU. |
| """ |
| idx = top_idx.permute(0, 3, 1, 2) |
| w = top_weight.permute(0, 3, 1, 2) |
| gathered_logits = torch.gather(logits, 1, idx) |
| logsumexp = torch.logsumexp(logits, dim=1, keepdim=True) |
| gathered_logp = gathered_logits - logsumexp |
| per_pixel_nll = -(gathered_logp * w).sum(dim=1) |
|
|
| pixel_class_weight = class_weights[idx[:, 0]] |
| return (per_pixel_nll * pixel_class_weight).mean() |
|
|
|
|
| def lab_to_rgb_batch(L, ab): |
| """L:(N,1,H,W) normalized, ab:(N,2,H,W) raw Lab units -> (N,H,W,3) numpy RGB in [0,1].""" |
| L = (L.detach().cpu().numpy() + 1.0) * 50.0 |
| ab = ab.detach().cpu().numpy() |
| lab = np.concatenate([L, ab], axis=1).transpose(0, 2, 3, 1) |
| return np.stack([np.clip(lab2rgb(x), 0, 1) for x in lab]) |
|
|
|
|
| def save_sample_grid(model, L, ab_true, path, device, temperature=0.38): |
| model.eval() |
| with torch.no_grad(): |
| logits = model(L.to(device)) |
| ab_pred = model.decode(logits, temperature=temperature).cpu() |
| rgb_true = lab_to_rgb_batch(L, ab_true) |
| rgb_pred = lab_to_rgb_batch(L, ab_pred) |
| gray = ((L + 1.0) * 0.5).repeat(1, 3, 1, 1).clamp(0, 1) |
|
|
| n = min(6, L.shape[0]) |
| rows = [] |
| for i in range(n): |
| rows.append(gray[i]) |
| rows.append(torch.from_numpy(rgb_pred[i]).permute(2, 0, 1).float()) |
| rows.append(torch.from_numpy(rgb_true[i]).permute(2, 0, 1).float()) |
| save_image(make_grid(rows, nrow=3, padding=2), path) |
| model.train() |
|
|
|
|
| |
| |
| |
| def build_argparser(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--dataset", default="frgfm/imagenette") |
| p.add_argument("--dataset-config", default=None, |
| help="Config name, if the dataset needs one (frgfm/imagenette does: " |
| "'160px', '320px', or 'full_size'). Leave unset for single-config " |
| "datasets, e.g. --dataset johnowhitaker/imagenette2-320.") |
| p.add_argument("--image-size", type=int, default=128) |
| p.add_argument("--base-channels", type=int, default=44, |
| help="Backbone width. If --init-from is set, only layers whose shape " |
| "matches are warm-started -- so changing this vs. the checkpoint " |
| "means starting the whole backbone fresh, not just the head.") |
| p.add_argument("--init-from", default=None, |
| help="Hub model id (or local path) to warm-start weights from. Layers " |
| "whose shape matches (the encoder/decoder backbone) are loaded; " |
| "layers that don't (e.g. the output head, if the bin grid or " |
| "--base-channels differs) are left randomly initialized. Works " |
| "across --image-size changes since the model is fully convolutional.") |
| p.add_argument("--bin-grid-size", type=int, default=10, |
| help="Spacing (in Lab ab units) of the color-bin grid.") |
| p.add_argument("--bin-samples", type=int, default=3000, |
| help="Number of training images to sample when deriving the color-bin " |
| "set and class-rebalancing weights from the real data distribution.") |
| p.add_argument("--rebalance-lambda", type=float, default=0.5, |
| help="Mix-in of uniform weighting when computing class-rebalancing " |
| "weights (Zhang et al.'s lambda). 0 = pure inverse-frequency " |
| "(aggressively up-weights rare colors); 1 = uniform (no " |
| "rebalancing, closer to plain classification).") |
| p.add_argument("--annealed-temp", type=float, default=0.38, |
| help="Temperature for the annealed-mean decode used for sample-grid " |
| "previews. Lower = closer to the mode (more vivid, can be " |
| "blotchy); higher = closer to the mean (smoother, can be " |
| "desaturated again). 0.38 is Zhang et al.'s value.") |
| p.add_argument("--grayscale-chroma-threshold", type=float, default=3.0, |
| help="Mean Lab chroma below which a sampled image is treated as " |
| "genuinely grayscale-source (vs. a real but desaturated color " |
| "photo). Always applied when building color-bin statistics. Only " |
| "affects actual training data if --exclude-grayscale-source is set.") |
| p.add_argument("--exclude-grayscale-source", action="store_true", |
| help="Also drop grayscale-source images from the train/val sets " |
| "themselves, not just from the bin-statistics sample. Off by " |
| "default since it changes what the model is trained on -- worth " |
| "trying if a meaningful fraction of the dataset turns out to be " |
| "grayscale-source, since those examples contribute pure " |
| "'predict zero' gradient that works against the rebalancing.") |
| p.add_argument("--batch-size", type=int, default=64) |
| p.add_argument("--no-amp", action="store_true", |
| help="Disable mixed-precision training. AMP is on by default on CUDA -- " |
| "the (B,num_bins,H,W) logits tensor is the single largest activation " |
| "in this model (much bigger than the old regression head's B,2,H,W), " |
| "so halving its memory footprint via fp16 matters more here than in " |
| "most small models.") |
| p.add_argument("--epochs", type=int, default=60) |
| p.add_argument("--lr", type=float, default=2e-4) |
| p.add_argument("--head-lr", type=float, default=None, |
| help="Separate learning rate for the output head (out_conv). Defaults " |
| "to --lr when unset. Worth raising above --lr when --init-from is " |
| "used: the warm-started backbone benefits from a gentle fine-tuning " |
| "rate, but the classification head is *not* warm-started (its shape " |
| "depends on the bin grid, which usually differs from the checkpoint " |
| "you're resuming from) -- it starts from random weights and can be " |
| "undertrained if it only ever sees the backbone's cautious LR.") |
| p.add_argument("--context-lr", type=float, default=None, |
| help="Separate learning rate for the dilated-conv context block. " |
| "Defaults to --lr when unset. Worth raising when warm-starting from " |
| "a checkpoint that predates this block: unlike the backbone, it " |
| "starts randomly initialized, and its residual connection means it " |
| "starts as a near no-op -- at the backbone's cautious LR it may not " |
| "get enough signal to develop meaningful behavior within a modest " |
| "number of epochs.") |
| p.add_argument("--num-workers", type=int, default=min(8, os.cpu_count() or 4)) |
| p.add_argument("--output-dir", default="./output") |
| p.add_argument("--sample-every", type=int, default=5) |
| p.add_argument("--push-to-hub", action="store_true") |
| p.add_argument("--push-every", type=int, default=10, |
| help="Push a checkpoint to the Hub every N epochs, so a " |
| "job timeout only loses progress since the last " |
| "push instead of the whole run. Only used with " |
| "--push-to-hub.") |
| p.add_argument("--hub-model-id", default=None) |
| p.add_argument("--seed", type=int, default=0) |
| return p |
|
|
|
|
| def get_dataset_splits(dataset_name, dataset_config, seed=0, val_fraction=0.05): |
| from datasets import Image as HFImage |
| from datasets import load_dataset |
|
|
| print(f"Loading {dataset_name}" + (f" ({dataset_config})" if dataset_config else "") + " ...") |
| try: |
| ds = load_dataset(dataset_name, dataset_config) if dataset_config else load_dataset(dataset_name) |
| except RuntimeError as e: |
| if "Dataset scripts are no longer supported" not in str(e): |
| raise |
| |
| |
| |
| |
| |
| |
| |
| print("Dataset uses a legacy loading script; loading its Parquet " |
| f"files directly from refs/convert/parquet/{dataset_config}/ ...") |
| base = f"hf://datasets/{dataset_name}@refs%2Fconvert%2Fparquet/{dataset_config}" |
| try: |
| ds = load_dataset("parquet", data_files={ |
| "train": f"{base}/train/*.parquet", |
| "validation": f"{base}/validation/*.parquet", |
| }) |
| except ValueError as e2: |
| if "all the data_files are invalid" not in str(e2): |
| raise |
| raise RuntimeError( |
| f"No Parquet mirror found for {dataset_name} config '{dataset_config}' at " |
| f"refs/convert/parquet right now -- this can happen even for a config that " |
| f"worked before (transient Hub-side issue with this repo, not a bug here). " |
| f"Consider a dataset stored natively as Parquet instead (no legacy-script " |
| f"fallback needed at all), e.g. --dataset johnowhitaker/imagenette2-320." |
| ) from e2 |
| |
| |
| |
| ds = ds.cast_column("image", HFImage()) |
| train_split = ds["train"] |
| if "validation" in ds: |
| val_split = ds["validation"] |
| elif "test" in ds: |
| val_split = ds["test"] |
| else: |
| |
| |
| |
| |
| print(f"{dataset_name} has no validation/test split -- creating a " |
| f"{val_fraction:.0%} held-out split via train_test_split(seed={seed}).") |
| split = train_split.train_test_split(test_size=val_fraction, seed=seed) |
| train_split, val_split = split["train"], split["test"] |
| return train_split, val_split |
|
|
|
|
| def build_model_card(model, n_params, args, train_loss, val_loss, epoch, is_final): |
| status = (f"Final ({args.epochs} epochs complete)" if is_final |
| else f"Checkpoint at epoch {epoch}/{args.epochs} (training still in progress " |
| f"when this was pushed -- may be superseded by a later commit)") |
| return f"""--- |
| tags: |
| - image-to-image |
| - colorization |
| - unet |
| - pytorch |
| license: apache-2.0 |
| --- |
| |
| # Mini U-Net Colorizer |
| |
| A {n_params:,}-parameter U-Net that colorizes grayscale photos. Classification-style |
| (Zhang et al., "Colorful Image Colorization"): it predicts a distribution over |
| {model.num_bins} quantized CIE Lab `a`/`b` bins per pixel rather than regressing a |
| single ab value directly, with color-bin loss weights derived from the real |
| training-data distribution (rare/saturated colors weighted higher) so it |
| doesn't just hedge toward desaturated averages. Decode with an annealed mean. |
| |
| - **Status:** {status} |
| - **Input:** L channel normalized as `L/50 - 1` -> `[-1, 1]`, shape `(1, {args.image_size}, {args.image_size})` |
| - **Output:** logits over {model.num_bins} ab bins, shape `({model.num_bins}, {args.image_size}, {args.image_size})` |
| - **Trained on:** `{args.dataset}` ({args.dataset_config}){f", warm-started from `{args.init_from}`" if args.init_from else ""} |
| - **Loss so far (weighted soft cross-entropy):** train {train_loss:.4f}, val {val_loss:.4f} |
| |
| ## Usage |
| |
| ```python |
| import numpy as np, torch |
| from skimage.color import rgb2lab, lab2rgb |
| from PIL import Image |
| # paste the SmallUNetColorizer class definition from the training script, then: |
| model = SmallUNetColorizer.from_pretrained("{args.hub_model_id or 'USERNAME/REPO'}") |
| model.eval() |
| |
| img = Image.open("photo.jpg").convert("RGB").resize(({args.image_size}, {args.image_size})) |
| lab = rgb2lab(np.asarray(img).astype("float32") / 255.0) |
| L = torch.from_numpy(lab[:, :, 0:1] / 50.0 - 1.0).permute(2, 0, 1)[None] |
| |
| with torch.no_grad(): |
| logits = model(L) |
| ab = model.decode(logits, temperature=0.38)[0].permute(1, 2, 0).numpy() |
| |
| L_out = (L[0, 0].numpy() + 1) * 50.0 |
| lab_out = np.concatenate([L_out[:, :, None], ab], axis=-1) |
| rgb_out = np.clip(lab2rgb(lab_out), 0, 1) |
| ``` |
| """ |
|
|
|
|
| def push_model(model, out_dir, args, n_params, train_loss, val_loss, epoch, is_final): |
| """Push current weights + a model card to the Hub. Never raises -- a failed |
| push (e.g. transient network error) should not kill an otherwise-fine job, |
| especially mid-training where we'll just try again at the next checkpoint.""" |
| try: |
| model.push_to_hub(args.hub_model_id) |
| api = HfApi() |
| card = build_model_card(model, n_params, args, train_loss, val_loss, epoch, is_final) |
| api.upload_file(path_or_fileobj=card.encode(), path_in_repo="README.md", |
| repo_id=args.hub_model_id) |
| sample_files = sorted((out_dir / "samples").glob("*.png")) |
| if sample_files: |
| api.upload_file(path_or_fileobj=str(sample_files[-1]), path_in_repo="sample.png", |
| repo_id=args.hub_model_id) |
| tag = "final" if is_final else f"epoch {epoch}" |
| print(f"Pushed ({tag}) to https://huggingface.co/{args.hub_model_id}") |
| except Exception as e: |
| print(f"WARNING: push_to_hub failed at epoch {epoch} ({e}). " |
| f"Continuing -- local copy is still at {out_dir / 'model'}.") |
|
|
|
|
| def load_raw_state_dict(model_id_or_path): |
| """Load just the raw weight tensors from a checkpoint, without constructing |
| a model instance. This matters for warm-starting: the checkpoint we're |
| resuming from might be an older/incompatible architecture (e.g. a |
| previous regression-head version of this model) whose config.json can't |
| even be used to construct today's model class. We only need the tensors |
| themselves to warm-start by-name, shape-compatible layers -- not a |
| working instance of whatever architecture produced them.""" |
| p = Path(model_id_or_path) |
| if p.is_dir(): |
| safep = p / "model.safetensors" |
| if safep.exists(): |
| return load_safetensors(str(safep)) |
| return torch.load(p / "pytorch_model.bin", map_location="cpu") |
| try: |
| path = hf_hub_download(model_id_or_path, filename="model.safetensors") |
| return load_safetensors(path) |
| except Exception: |
| path = hf_hub_download(model_id_or_path, filename="pytorch_model.bin") |
| return torch.load(path, map_location="cpu") |
|
|
|
|
| def filter_grayscale_source(hf_split, threshold): |
| """Drop images whose mean chroma is below `threshold` -- genuinely |
| grayscale-source photos, which contribute pure 'predict zero' gradient |
| with no offsetting colorful signal, working against the very thing the |
| rebalancing weights are trying to achieve.""" |
| def is_colorful(example): |
| img = example["image"] |
| if img.mode != "RGB": |
| img = img.convert("RGB") |
| img = img.resize((64, 64)) |
| arr = np.asarray(img).astype(np.float32) / 255.0 |
| return image_mean_chroma(rgb2lab(arr)) >= threshold |
|
|
| n_before = len(hf_split) |
| filtered = hf_split.filter(is_colorful, num_proc=min(8, os.cpu_count() or 1)) |
| n_after = len(filtered) |
| print(f"Excluded {n_before - n_after}/{n_before} grayscale-source images " |
| f"(mean chroma < {threshold}) from this split.") |
| return filtered |
|
|
|
|
| def main(argv=None, dataset_loader=get_dataset_splits): |
| args = build_argparser().parse_args(argv) |
| if args.push_to_hub and not args.hub_model_id: |
| raise SystemExit("--hub-model-id is required with --push-to-hub") |
|
|
| torch.manual_seed(args.seed) |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| print(f"Device: {device}") |
| out_dir = Path(args.output_dir) |
| (out_dir / "samples").mkdir(parents=True, exist_ok=True) |
|
|
| train_split, val_split = dataset_loader(args.dataset, args.dataset_config, seed=args.seed) |
|
|
| if args.exclude_grayscale_source: |
| print("Filtering grayscale-source images out of train/val (--exclude-grayscale-source)...") |
| train_split = filter_grayscale_source(train_split, args.grayscale_chroma_threshold) |
| val_split = filter_grayscale_source(val_split, args.grayscale_chroma_threshold) |
|
|
| print(f"Deriving color bins from {args.bin_samples} sampled training images " |
| f"(grid={args.bin_grid_size}, lambda={args.rebalance_lambda}) ...") |
| color_bins = build_color_bins(train_split, num_samples=args.bin_samples, |
| grid_size=args.bin_grid_size, |
| rebalance_lambda=args.rebalance_lambda, seed=args.seed, |
| grayscale_chroma_threshold=args.grayscale_chroma_threshold) |
| print(f"{color_bins.num_bins} color bins in use " |
| f"(weight range {color_bins.weights.min():.2f}-{color_bins.weights.max():.2f})") |
|
|
| model = SmallUNetColorizer(bin_centers=color_bins.bin_centers.tolist(), |
| base=args.base_channels).to(device) |
| n_params = sum(x.numel() for x in model.parameters()) |
| print(f"Model parameters: {n_params:,}") |
| assert n_params < 4_000_000, f"Model has {n_params:,} params, over the 4M budget!" |
|
|
| if args.init_from: |
| print(f"Warm-starting from {args.init_from} ...") |
| ckpt_sd = load_raw_state_dict(args.init_from) |
| own_sd = model.state_dict() |
| compatible = {k: v for k, v in ckpt_sd.items() |
| if k in own_sd and own_sd[k].shape == v.shape} |
| skipped = sorted(set(own_sd) - set(compatible)) |
| model.load_state_dict(compatible, strict=False) |
| print(f"Loaded {len(compatible)}/{len(own_sd)} tensors from the checkpoint " |
| f"(shape-compatible backbone layers). Randomly initialized: {skipped or 'none'}.") |
|
|
| class_weights = torch.from_numpy(color_bins.weights).to(device) |
|
|
| train_ds = ColorizationDataset(train_split, color_bins, args.image_size, augment=True) |
| val_ds = ColorizationDataset(val_split, color_bins, args.image_size, augment=False) |
| print(f"train: {len(train_ds)} images, val: {len(val_ds)} images") |
|
|
| common_kwargs = dict(num_workers=args.num_workers, pin_memory=(device == "cuda")) |
| if args.num_workers > 0: |
| common_kwargs["persistent_workers"] = True |
| train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, |
| drop_last=True, **common_kwargs) |
| val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, **common_kwargs) |
|
|
| |
| |
| |
| |
| |
| |
| sample_rng = np.random.default_rng(args.seed) |
| n_preview = min(args.batch_size, len(val_ds)) |
| preview_indices = sample_rng.choice(len(val_ds), size=n_preview, replace=False).tolist() |
| preview_loader = DataLoader(torch.utils.data.Subset(val_ds, preview_indices), |
| batch_size=n_preview, shuffle=False) |
| fixed_val_L, fixed_val_ab, _, _ = next(iter(preview_loader)) |
|
|
| head_lr = args.head_lr if args.head_lr is not None else args.lr |
| context_lr = args.context_lr if args.context_lr is not None else args.lr |
| backbone_params = [p for n, p in model.named_parameters() |
| if not n.startswith("out_conv") and not n.startswith("context")] |
| context_params = [p for n, p in model.named_parameters() if n.startswith("context")] |
| head_params = [p for n, p in model.named_parameters() if n.startswith("out_conv")] |
| opt = torch.optim.AdamW([ |
| {"params": backbone_params, "lr": args.lr}, |
| {"params": context_params, "lr": context_lr}, |
| {"params": head_params, "lr": head_lr}, |
| ], weight_decay=1e-4) |
| print(f"Learning rates: backbone={args.lr} " |
| f"context={context_lr}{' (== --lr)' if args.context_lr is None else ''} " |
| f"head={head_lr}{' (== --lr)' if args.head_lr is None else ''}") |
| sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs) |
| amp_enabled = (device == "cuda") and not args.no_amp |
| scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled) |
| print(f"Mixed precision (AMP): {'on' if amp_enabled else 'off'}") |
|
|
| def criterion(logits, top_idx, top_weight): |
| return classification_loss(logits, top_idx, top_weight, class_weights) |
|
|
| train_loss = val_loss = float("nan") |
| for epoch in range(1, args.epochs + 1): |
| epoch_start = time.time() |
| model.train() |
| running = 0.0 |
| pbar = tqdm(train_loader, desc=f"epoch {epoch}/{args.epochs}") |
| for L, ab, top_idx, top_weight in pbar: |
| L = L.to(device, non_blocking=True) |
| top_idx = top_idx.to(device, non_blocking=True) |
| top_weight = top_weight.to(device, non_blocking=True) |
| with torch.amp.autocast("cuda", enabled=amp_enabled): |
| logits = model(L) |
| loss = criterion(logits, top_idx, top_weight) |
| opt.zero_grad(set_to_none=True) |
| scaler.scale(loss).backward() |
| scaler.step(opt) |
| scaler.update() |
| running += loss.item() * L.size(0) |
| pbar.set_postfix(loss=f"{loss.item():.4f}") |
| sched.step() |
| train_loss = running / len(train_ds) |
|
|
| model.eval() |
| val_running = 0.0 |
| with torch.no_grad(), torch.amp.autocast("cuda", enabled=amp_enabled): |
| for L, ab, top_idx, top_weight in val_loader: |
| L = L.to(device) |
| top_idx, top_weight = top_idx.to(device), top_weight.to(device) |
| val_running += criterion(model(L), top_idx, top_weight).item() * L.size(0) |
| val_loss = val_running / len(val_ds) |
| print(f"epoch {epoch}: train_loss={train_loss:.4f} val_loss={val_loss:.4f} " |
| f"({time.time() - epoch_start:.1f}s)") |
|
|
| if epoch % args.sample_every == 0 or epoch == args.epochs: |
| save_sample_grid(model, fixed_val_L, fixed_val_ab, |
| out_dir / "samples" / f"epoch_{epoch:03d}.png", device, |
| temperature=args.annealed_temp) |
|
|
| if args.push_to_hub and epoch != args.epochs and epoch % args.push_every == 0: |
| push_model(model, out_dir, args, n_params, train_loss, val_loss, epoch, is_final=False) |
|
|
| model_dir = out_dir / "model" |
| model.save_pretrained(model_dir) |
| print(f"Saved model to {model_dir}") |
|
|
| if args.push_to_hub: |
| push_model(model, out_dir, args, n_params, train_loss, val_loss, args.epochs, is_final=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
