# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch>=2.3",
#     "torchvision>=0.18",
#     "datasets>=2.19",
#     "huggingface_hub>=0.24",
#     "safetensors>=0.4",
#     "scikit-image>=0.22",
#     "scipy>=1.11",
#     "pillow>=10.0",
#     "numpy",
#     "tqdm",
# ]
# ///
"""
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


# --------------------------------------------------------------------------
# Classification-based colorization (Zhang et al., "Colorful Image
# Colorization" style): quantize ab space into bins, predict a distribution
# over bins per pixel instead of regressing a single ab value, and decode
# with an annealed mean. This avoids the "regress to the desaturated average
# of all plausible colors" failure mode of plain L1/L2 regression.
# --------------------------------------------------------------------------
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)  # (Q, 2)
        self.weights = weights.astype(np.float32)  # (Q,)
        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)  # (N,k), (N,k)
        if k == 1:  # cKDTree squeezes the last dim when 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  # bin centers along one axis
    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))  # small is plenty for a color histogram
        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)  # (Q, 2)

    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()  # normalize so E_prior[w] == 1

    return ColorBins(bin_centers, w)


# --------------------------------------------------------------------------
# Model: small U-Net, L (1ch) in -> per-pixel logits over color bins out.
# ~3.6M trainable params at base=44 (bin_centers is a non-trainable buffer).
# --------------------------------------------------------------------------
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)  # bottleneck
        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)  # raw logits, no activation

    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)  # (B, Q, H, W) logits

    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)  # (B,Q,H,W)
        return torch.einsum("bqhw,qc->bchw", probs_t, self.bin_centers)


# --------------------------------------------------------------------------
# Data: RGB photo -> (L, ab) in normalized Lab space
# --------------------------------------------------------------------------
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]  # raw Lab units, kept continuous for viz ground truth

        top_idx, top_weight = self.color_bins.soft_encode(ab)  # (H,W,5) each

        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)      # (H,W,5) int64
        top_weight = torch.from_numpy(top_weight)  # (H,W,5) float32
        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)      # (B,5,H,W)
    w = top_weight.permute(0, 3, 1, 2)     # (B,5,H,W)
    gathered_logits = torch.gather(logits, 1, idx)         # (B,5,H,W) -- small
    logsumexp = torch.logsumexp(logits, dim=1, keepdim=True)  # (B,1,H,W) -- small
    gathered_logp = gathered_logits - logsumexp            # (B,5,H,W)
    per_pixel_nll = -(gathered_logp * w).sum(dim=1)  # (B,H,W)

    pixel_class_weight = class_weights[idx[:, 0]]  # nearest bin per pixel, (B,H,W)
    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()


# --------------------------------------------------------------------------
# Training
# --------------------------------------------------------------------------
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
        # This repo only ships a legacy loading script, which `datasets` now
        # refuses to run. It still has an auto-converted Parquet mirror on
        # `refs/convert/parquet`, but for repos with multiple size configs
        # (like this one: full_size/160px/320px) that mirror doesn't expose
        # them as separate named configs -- so instead of going through
        # `datasets`' config auto-detection, load the underlying Parquet
        # shards for this config directly.
        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
        # The generic "parquet" builder doesn't always restore the
        # decode-on-access Image feature -- make sure it's actually an
        # Image column and not raw {"bytes":..,"path":..} dicts.
        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:
        # No separate held-out split at all (e.g. johnowhitaker/imagenette2-320
        # ships everything as a single "train"). Carve out a genuine held-out
        # set ourselves rather than reusing training data for validation --
        # val_loss and the sample-grid preview should reflect unseen images.
        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)

    # For the recurring sample-grid preview, draw a random (but seeded, so
    # reproducible run-to-run) sample spanning the whole validation set --
    # not just the literal first batch of an unshuffled loader, which for a
    # class-grouped dataset like imagenette means a single class every time.
    # (val_loss above is unaffected either way -- it already iterates the
    # full val_loader, all classes, every epoch.)
    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()