"""Self-contained CIFAR-10 classifier loader for the deployed models. No dependency on the training repo — only torch, torchvision, timm and huggingface_hub. Works with a local model folder or a Hugging Face Hub repo id. Each model folder contains: - config.json model + preprocessing metadata - model.safetensors weights (or pytorch_model.bin as fallback) - README.md model card Example: from cifar_classifier import CIFAR10Classifier clf = CIFAR10Classifier.from_pretrained("YOUR_USER/convnextv2-huge-cifar10-upsample") label, probs = clf.predict("cat.jpg") """ from __future__ import annotations import json import os from pathlib import Path from typing import Any import torch import torch.nn as nn import timm from PIL import Image from torchvision import transforms CIFAR10_CLASSES = [ "airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck", ] _WEIGHT_FILES = ("model.safetensors", "pytorch_model.bin") def _replace_conv(conv: nn.Conv2d, stride: int = 1) -> nn.Conv2d: """3x3 stride-1 stem conv used for native 32x32 CIFAR input. Mirrors ``adapt_stem_for_cifar`` in the training repo (src/models/registry.py). Weights are loaded from the checkpoint, so the averaged-kernel initialisation used at training time is irrelevant here — only the shapes must match. """ return nn.Conv2d( conv.in_channels, conv.out_channels, kernel_size=3, stride=stride, padding=1, bias=conv.bias is not None, ) def adapt_stem_for_cifar(model: nn.Module, model_name: str) -> None: """Adapt timm stems for native 32x32 CIFAR input (must match training).""" if hasattr(model, "patch_embed") and hasattr(model.patch_embed, "proj"): model.patch_embed.proj = _replace_conv(model.patch_embed.proj) return if hasattr(model, "stem"): stem = model.stem if hasattr(stem, "conv") and isinstance(stem.conv, nn.Conv2d): stem.conv = _replace_conv(stem.conv) elif isinstance(stem, nn.Sequential): for module in stem.modules(): if isinstance(module, nn.Conv2d): module.kernel_size = (3, 3) module.stride = (1, 1) module.padding = (1, 1) break if model_name.startswith("convnext") and hasattr(model, "downsample_layers"): first = model.downsample_layers[0] if isinstance(first, nn.Sequential) and isinstance(first[0], nn.Conv2d): first[0] = _replace_conv(first[0]) if model_name.startswith("tf_efficientnet") and hasattr(model, "conv_stem"): model.conv_stem = _replace_conv(model.conv_stem) def build_model(config: dict[str, Any]) -> nn.Module: """Build the timm model exactly as configured at training time.""" model = timm.create_model( config["timm_name"], pretrained=False, num_classes=config.get("num_classes", 10), ) if config.get("input_mode", "upsample") == "native": adapt_stem_for_cifar(model, config["timm_name"]) return model def build_eval_transform(config: dict[str, Any]) -> transforms.Compose: """Deterministic preprocessing matching the training eval pipeline.""" ops: list[Any] = [] size = config.get("image_size", 32) if config.get("input_mode", "upsample") == "upsample": ops.append(transforms.Resize((size, size))) else: # Native models were trained on 32x32; downscale larger inputs. ops.append(transforms.Resize((size, size))) ops.append(transforms.ToTensor()) ops.append(transforms.Normalize(mean=config["mean"], std=config["std"])) return transforms.Compose(ops) def _load_state_dict(path: Path) -> dict[str, torch.Tensor]: if path.suffix == ".safetensors": from safetensors.torch import load_file return load_file(str(path)) return torch.load(path, map_location="cpu", weights_only=True) class CIFAR10Classifier: """Inference wrapper around a deployed CIFAR-10 model folder.""" def __init__(self, model: nn.Module, config: dict[str, Any], device: str | torch.device = "cpu"): self.config = config self.classes: list[str] = config.get("classes", CIFAR10_CLASSES) self.device = torch.device(device) self.model = model.to(self.device).eval() self.transform = build_eval_transform(config) # -- loading --------------------------------------------------------- @classmethod def from_pretrained( cls, repo_id_or_path: str | os.PathLike[str], device: str | torch.device | None = None, revision: str | None = None, ) -> "CIFAR10Classifier": """Load from a local folder or a Hugging Face Hub repo id.""" folder = cls._resolve_folder(repo_id_or_path, revision) with open(folder / "config.json", encoding="utf-8") as fh: config = json.load(fh) model = build_model(config) for name in _WEIGHT_FILES: weight_path = folder / name if weight_path.exists(): model.load_state_dict(_load_state_dict(weight_path)) break else: raise FileNotFoundError(f"No weight file ({', '.join(_WEIGHT_FILES)}) found in {folder}") if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" return cls(model, config, device) @staticmethod def _resolve_folder(repo_id_or_path: str | os.PathLike[str], revision: str | None) -> Path: path = Path(repo_id_or_path) if path.is_dir(): return path from huggingface_hub import snapshot_download return Path( snapshot_download( repo_id=str(repo_id_or_path), revision=revision, allow_patterns=["config.json", "README.md", *_WEIGHT_FILES], ) ) # -- inference --------------------------------------------------------- def _preprocess(self, image: str | os.PathLike[str] | Image.Image) -> torch.Tensor: if not isinstance(image, Image.Image): image = Image.open(image) image = image.convert("RGB") return self.transform(image) @torch.no_grad() def predict(self, image: str | os.PathLike[str] | Image.Image) -> tuple[str, dict[str, float]]: """Return (top label, {class: probability}) for one image.""" tensor = self._preprocess(image).unsqueeze(0).to(self.device) probs = torch.softmax(self.model(tensor), dim=-1)[0].cpu() top = int(probs.argmax()) return self.classes[top], {cls: float(probs[i]) for i, cls in enumerate(self.classes)} @torch.no_grad() def predict_batch(self, images: list[str | os.PathLike[str] | Image.Image]) -> list[str]: """Return the top label for each image.""" batch = torch.stack([self._preprocess(img) for img in images]).to(self.device) indices = self.model(batch).argmax(dim=-1).cpu().tolist() return [self.classes[i] for i in indices] if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Run CIFAR-10 inference with a deployed model.") parser.add_argument("model", help="Local model folder or Hub repo id (e.g. user/convnextv2-huge-cifar10-upsample)") parser.add_argument("image", help="Path to an image file") parser.add_argument("--device", default=None, help="cpu | cuda (default: auto)") args = parser.parse_args() clf = CIFAR10Classifier.from_pretrained(args.model, device=args.device) label, probs = clf.predict(args.image) print(f"Prediction: {label}") for cls_name, p in sorted(probs.items(), key=lambda kv: kv[1], reverse=True): print(f" {cls_name:<12s} {p:.4f}")