Upload 3 files

colorize_eval 7.py

@@ -0,0 +1,357 @@
+# /// 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",
+#     "pillow>=10.0",
+#     "numpy",
+# ]
+# ///
+"""
+Evaluate a SmallUNetColorizer checkpoint with a large, class-stratified
+sample grid -- unlike the small fixed grid saved during training (which,
+before a fix, was drawn from a single class), this samples N images from
+*every* class in the dataset, labeled, so weaknesses specific to one subject
+category (e.g. reflective electronics) are visible instead of hidden.
+
+Runs entirely on CPU -- inference-only on a ~3.6M param model, no GPU needed.
+Designed for `hf jobs uv run --flavor cpu-basic`; since a Jobs container is
+destroyed when it finishes, the resulting grid is pushed to a Hub repo
+rather than saved somewhere you could copy it from afterward.
+
+    hf jobs uv run --flavor cpu-basic --timeout 1h -s HF_TOKEN \\
+        colorize_eval.py \\
+        --model User-2468/mini-unet-colorizer \\
+        --per-class 5 --temperature 0.38 \\
+        --push-to-hub --hub-model-id User-2468/mini-unet-colorizer
+"""
+import argparse
+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
+from PIL import Image, ImageDraw, ImageFont
+from skimage.color import lab2rgb, rgb2lab
+
+
+# --------------------------------------------------------------------------
+# Model (identical to colorize_train.py -- kept self-contained since Jobs
+# scripts don't have access to sibling files).
+# --------------------------------------------------------------------------
+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),
+        )
+        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):
+        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):
+        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)
+
+
+# --------------------------------------------------------------------------
+# Dataset loading (identical fallback logic to colorize_train.py, for
+# datasets like frgfm/imagenette that ship a legacy loading script).
+# --------------------------------------------------------------------------
+# Some Hub repos for imagenette label classes by WordNet synset ID rather
+# than a readable name (e.g. johnowhitaker/imagenette2-320). Standard,
+# well-documented imagenette synset -> name mapping so captions stay
+# readable regardless of which repo's labeling convention is in use.
+SYNSET_TO_NAME = {
+    "n01440764": "tench", "n02102040": "English springer",
+    "n02979186": "cassette player", "n03000684": "chain saw",
+    "n03028079": "church", "n03394916": "French horn",
+    "n03417042": "garbage truck", "n03425413": "gas pump",
+    "n03445777": "golf ball", "n03888257": "parachute",
+}
+
+
+def get_dataset_splits(dataset_name, dataset_config):
+    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:
+        eval_split = ds["validation"]
+    elif "test" in ds:
+        eval_split = ds["test"]
+    else:
+        # Some repos (e.g. johnowhitaker/imagenette2-320) ship a single
+        # unsplit "train" -- fine for a qualitative visual check like this
+        # script does, just note it's not a strictly held-out set.
+        print(f"Note: {dataset_name} has no validation/test split -- sampling "
+              f"from 'train' instead. Fine for a visual check, but these may "
+              f"overlap with what the model was actually trained on.")
+        eval_split = train_split
+    return train_split, eval_split
+
+
+# --------------------------------------------------------------------------
+# Evaluation
+# --------------------------------------------------------------------------
+def colorize_one(model, img, size, temperature, grayscale_chroma_threshold=3.0):
+    """Returns (gray_rgb, pred_rgb, true_rgb, is_grayscale_source)."""
+    img = img.convert("RGB").resize((size, size))
+    arr = np.asarray(img).astype(np.float32) / 255.0
+    lab = rgb2lab(arr).astype(np.float32)
+    L = torch.from_numpy(lab[:, :, 0:1] / 50.0 - 1.0).permute(2, 0, 1)[None]
+
+    with torch.no_grad():
+        logits = model(L)
+        ab_pred = model.decode(logits, temperature=temperature)[0].permute(1, 2, 0).numpy()
+
+    L_raw = lab[:, :, 0]
+    ab_true = lab[:, :, 1:3]
+    is_grayscale = float(np.sqrt((ab_true ** 2).sum(-1)).mean()) < grayscale_chroma_threshold
+
+    def to_rgb(ab):
+        lab_full = np.concatenate([L_raw[:, :, None], ab], axis=-1)
+        return (np.clip(lab2rgb(lab_full), 0, 1) * 255).astype(np.uint8)
+
+    gray_rgb = (np.stack([L_raw] * 3, axis=-1) / 100.0 * 255).astype(np.uint8)
+    return gray_rgb, to_rgb(ab_pred), to_rgb(ab_true), is_grayscale
+
+
+def build_stratified_sample(val_split, per_class, seed):
+    """Returns list of (index, class_name), `per_class` images from every
+    class in the dataset -- not just whichever class happens to sort first."""
+    labels = val_split["label"]
+    raw_names = val_split.features["label"].names
+    class_names = [SYNSET_TO_NAME.get(n, n) for n in raw_names]  # readable, if known
+    rng = np.random.default_rng(seed)
+
+    by_class = {c: [] for c in range(len(class_names))}
+    for i, lbl in enumerate(labels):
+        by_class[lbl].append(i)
+
+    sample = []
+    for c, indices in by_class.items():
+        chosen = rng.choice(indices, size=min(per_class, len(indices)), replace=False)
+        sample.extend((int(i), class_names[c]) for i in chosen)
+    return sample
+
+
+def build_grid_image(rows_by_class, thumb_size, per_class):
+    """rows_by_class: dict class_name -> list of (gray, pred, true, is_grayscale_source)."""
+    font = ImageFont.load_default()
+    header_h = 24
+    pad = 2
+    row_w = thumb_size * 3 + pad * 2
+    class_block_h = header_h + per_class * (thumb_size + pad)
+
+    class_names = list(rows_by_class.keys())
+    total_h = class_block_h * len(class_names)
+    canvas = Image.new("RGB", (row_w, total_h), "white")
+    draw = ImageDraw.Draw(canvas)
+
+    y = 0
+    for name in class_names:
+        rows = rows_by_class[name]
+        n_gray_source = sum(1 for r in rows if r[3])
+        header_text = f"{name}  (gray | prediction | ground truth)"
+        if n_gray_source:
+            header_text += f"  -- {n_gray_source}/{len(rows)} grayscale-source"
+        draw.rectangle([0, y, row_w, y + header_h], fill=(30, 30, 30))
+        draw.text((6, y + 5), header_text, fill="white", font=font)
+        y += header_h
+        for gray, pred, true, is_gray_source in rows:
+            x = 0
+            for arr in (gray, pred, true):
+                im = Image.fromarray(arr).resize((thumb_size, thumb_size))
+                canvas.paste(im, (x, y))
+                x += thumb_size + pad
+            if is_gray_source:
+                # Tag the ground-truth thumbnail so it's obvious *which* row
+                # had nothing to colorize, not just that the class has some.
+                tag_x = thumb_size * 2 + pad * 2
+                draw.rectangle([tag_x, y, tag_x + thumb_size, y + 14], fill=(200, 30, 30))
+                draw.text((tag_x + 3, y + 2), "B&W source", fill="white", font=font)
+            y += thumb_size + pad
+    return canvas
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--model", required=True, help="Hub model id or local checkpoint path")
+    p.add_argument("--dataset", default="johnowhitaker/imagenette2-320",
+                    help="Default is stored natively as Parquet (no legacy-script fallback "
+                         "needed), which has been more reliable than frgfm/imagenette's "
+                         "auto-converted mirror. It has no separate validation split (just "
+                         "'train'), so this samples from the full set rather than a held-out "
+                         "portion -- fine for a visual check. frgfm/imagenette is still an "
+                         "option (--dataset-config 160px or 320px) if its mirror comes back.")
+    p.add_argument("--dataset-config", default=None,
+                    help="Config name, if the dataset needs one (frgfm/imagenette does: "
+                         "'160px' or '320px'). Leave unset for single-config datasets like "
+                         "the johnowhitaker default.")
+    p.add_argument("--image-size", type=int, default=256, help="Resolution used for inference")
+    p.add_argument("--thumb-size", type=int, default=128, help="Display size in the grid")
+    p.add_argument("--per-class", type=int, default=5, help="Sampled images per class")
+    p.add_argument("--temperature", type=float, default=0.38)
+    p.add_argument("--grayscale-chroma-threshold", type=float, default=3.0,
+                    help="Mean Lab chroma below which a ground-truth image is flagged as "
+                         "genuinely grayscale-source and tagged in the grid, rather than a "
+                         "real colorization miss.")
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--output-dir", default="./eval")
+    p.add_argument("--output-name", default="eval_grid.png")
+    p.add_argument("--push-to-hub", action="store_true")
+    p.add_argument("--hub-model-id", default=None,
+                    help="Repo to push the grid to (defaults to --model). Since a Jobs "
+                         "container is destroyed on completion, --push-to-hub is the only "
+                         "way to get the result back out when running as a Job.")
+    args = p.parse_args()
+
+    print(f"Loading {args.model} ...")
+    model = SmallUNetColorizer.from_pretrained(args.model).eval()
+    print(f"({model.num_bins} color bins)")
+
+    _, val_split = get_dataset_splits(args.dataset, args.dataset_config)
+    sample = build_stratified_sample(val_split, args.per_class, args.seed)
+    class_names = [SYNSET_TO_NAME.get(n, n) for n in val_split.features["label"].names]
+    print(f"Evaluating {len(sample)} images across {len(class_names)} classes "
+          f"({args.per_class} per class) ...")
+
+    rows_by_class = {name: [] for name in class_names}
+    n_grayscale_total = 0
+    for n, (idx, class_name) in enumerate(sample):
+        img = val_split[idx]["image"]
+        gray, pred, true, is_gray_source = colorize_one(
+            model, img, args.image_size, args.temperature, args.grayscale_chroma_threshold)
+        rows_by_class[class_name].append((gray, pred, true, is_gray_source))
+        n_grayscale_total += is_gray_source
+        if (n + 1) % 10 == 0 or n + 1 == len(sample):
+            print(f"  {n + 1}/{len(sample)}")
+
+    print(f"{n_grayscale_total}/{len(sample)} sampled images are grayscale-source "
+          f"(tagged in red in the grid) -- these have no color to recover, so weak-looking "
+          f"results there aren't a model failure.")
+
+    grid = build_grid_image(rows_by_class, args.thumb_size, args.per_class)
+
+    out_dir = Path(args.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out_path = out_dir / args.output_name
+    grid.save(out_path)
+    print(f"Saved grid ({grid.width}x{grid.height}) to {out_path}")
+
+    if args.push_to_hub:
+        hub_model_id = args.hub_model_id or args.model
+        try:
+            HfApi().upload_file(path_or_fileobj=str(out_path), path_in_repo=args.output_name,
+                                 repo_id=hub_model_id)
+            print(f"Pushed to https://huggingface.co/{hub_model_id}/blob/main/{args.output_name}")
+        except Exception as e:
+            print(f"WARNING: push failed ({e}). Grid was still generated at {out_path}, "
+                  f"but that path won't survive a Jobs container exiting.")
+
+
+if __name__ == "__main__":
+    main()

colorize_image 3.py

@@ -0,0 +1,176 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "torch>=2.3",
+#     "torchvision>=0.18",
+#     "huggingface_hub>=0.24",
+#     "safetensors>=0.4",
+#     "scikit-image>=0.22",
+#     "pillow>=10.0",
+#     "numpy",
+# ]
+# ///
+"""
+Colorize photos with a SmallUNetColorizer checkpoint trained by colorize_train.py
+(classification-head version: predicts a distribution over quantized Lab ab
+bins per pixel, decoded with an annealed mean).
+
+Usage:
+    uv run colorize_image.py --model User-2468/mini-unet-colorizer photo.jpg
+    uv run colorize_image.py --model ./local_checkpoint --temperature 0.2 a.jpg b.jpg
+
+--temperature is the main "vividness" knob: lower values weight the decode
+toward the most likely color bin (more saturated, can be a bit blotchy);
+higher values move toward the full expectation over the distribution
+(smoother, but can drift back toward desaturated -- the same hedging effect
+plain regression had). 0.38 (the default) is a reasonable middle ground.
+--saturation-boost is an optional *additional* post-decode multiplier on top
+of that, for further hand-tuning after picking a temperature.
+"""
+import argparse
+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 PyTorchModelHubMixin
+from PIL import Image
+from skimage.color import lab2rgb, rgb2lab
+
+
+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),
+        )
+        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):
+        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):
+        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)
+
+
+def colorize(model, img, size, temperature, saturation_boost, device):
+    img = img.convert("RGB").resize((size, size))
+    arr = np.asarray(img).astype(np.float32) / 255.0
+    lab = rgb2lab(arr).astype(np.float32)
+    L = torch.from_numpy(lab[:, :, 0:1] / 50.0 - 1.0).permute(2, 0, 1)[None].to(device)
+
+    with torch.no_grad():
+        logits = model(L)
+        ab = model.decode(logits, temperature=temperature)[0].permute(1, 2, 0).cpu().numpy()
+
+    ab = np.clip(ab * saturation_boost, -128, 127)
+    L_out = (L[0, 0].cpu().numpy() + 1.0) * 50.0
+    lab_out = np.concatenate([L_out[:, :, None], ab], axis=-1)
+    rgb_out = np.clip(lab2rgb(lab_out), 0, 1)
+    return Image.fromarray((rgb_out * 255).astype(np.uint8))
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("images", nargs="+", help="Path(s) to input photo(s)")
+    p.add_argument("--model", required=True, help="Hub model id or local checkpoint path")
+    p.add_argument("--size", type=int, default=256, help="Resize input to this square size")
+    p.add_argument("--temperature", type=float, default=0.38,
+                    help="Annealed-mean decode temperature. Lower = more vivid/mode-like, "
+                         "higher = smoother but can desaturate again. Try 0.15-0.6.")
+    p.add_argument("--saturation-boost", type=float, default=1.0,
+                    help="Extra multiplier on the decoded ab, applied after temperature. "
+                         "1.0 = no extra boost.")
+    p.add_argument("--output-dir", default="./colorized")
+    args = p.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Loading {args.model} on {device} ...")
+    model = SmallUNetColorizer.from_pretrained(args.model).to(device).eval()
+    print(f"({model.num_bins} color bins)")
+
+    out_dir = Path(args.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    for path in args.images:
+        img = Image.open(path)
+        result = colorize(model, img, args.size, args.temperature, args.saturation_boost, device)
+        out_path = out_dir / f"{Path(path).stem}_colorized.png"
+        result.save(out_path)
+        print(f"{path} -> {out_path}")
+
+
+if __name__ == "__main__":
+    main()

colorize_train 14.py

@@ -0,0 +1,786 @@
+# /// 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()