Add merged model + processor

config.json

@@ -1,7 +1,7 @@
 {
   "apply_layernorm": true,
   "architectures": [
-    "FontClassifierWithPreprocessing"
+    "Dinov2ForImageClassification"
   ],
   "attention_probs_dropout_prob": 0.0,
   "drop_path_rate": 0.0,
@@ -837,8 +837,5 @@
   "torch_dtype": "float32",
   "transformers_version": "4.52.4",
   "use_mask_token": true,
-  "use_swiglu_ffn": false,
-  "auto_map": {
-    "AutoModelForImageClassification": "font_classifier_with_preprocessing.FontClassifierWithPreprocessing"
-  }
-}
+  "use_swiglu_ffn": false
+}

handler.py

@@ -0,0 +1,63 @@
+# to be bundled with the model on upload to HF Inference Endpoints
+
+import base64
+import io
+from typing import Any, Dict
+
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, Dinov2ForImageClassification
+
+from train_model import get_inference_transform
+
+
+class EndpointHandler:
+    """
+    HF Inference Endpoints entry‑point.
+    Loads model/processor once, then uses your *imported* preprocessing
+    on every request.
+    """
+
+    def __init__(self, path: str = "", image_size: int = 224):
+        # Weights + processor --------------------------------------------------------
+        self.processor = AutoImageProcessor.from_pretrained(path or ".")
+        self.model     = (
+            Dinov2ForImageClassification.from_pretrained(path or ".")
+            .eval()
+        )
+
+        # Re‑use the exact transform from your code ---------------------------------
+        self.transform = get_inference_transform(self.processor, image_size)
+
+        self.id2label = self.model.config.id2label
+
+    # -------------------------------------------------------------------------------
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Expects {"inputs": "<base64‑encoded image>"}.
+        Returns the top prediction + per‑class probabilities.
+        """
+        if "inputs" not in data:
+            raise ValueError("Request JSON must contain an 'inputs' field.")
+
+        # Decode base64 → PIL
+        img_bytes = base64.b64decode(data["inputs"])
+        image = Image.open(io.BytesIO(img_bytes))
+
+        # Preprocess with your own transform
+        pixel_values = self.transform(image).unsqueeze(0)   # [1, C, H, W]
+
+        with torch.no_grad():
+            logits = self.model(pixel_values).logits
+            probs  = logits.softmax(dim=-1)[0]
+
+        top_idx = int(probs.argmax())
+        top_label = self.id2label[top_idx]
+
+        return {
+            "predicted_label": top_label,
+            "scores": {
+                self.id2label[i]: float(p)
+                for i, p in enumerate(probs)
+            }
+        }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0eeabb74d1af47629e61d6d4dd48bbf3eb74121db29c8ba8b644b41b8c481a6d
-size 348770168
+oid sha256:ecb73dab2fc1203ab36f177c7a6c5a23e472f5fff58b8ce5f8fc51f20f0480e1
+size 348769976

requirements.txt

@@ -0,0 +1,2 @@
+torchvision>=0.19
+Pillow>=10

train_model.py

@@ -0,0 +1,347 @@
+import argparse
+import json
+import logging
+import os
+import shutil
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchvision.transforms as T
+from datasets import load_dataset
+from huggingface_hub import HfApi
+from peft import LoraConfig, PeftModel, get_peft_model
+from PIL import Image
+from safetensors import safe_open
+from transformers import (
+    AutoImageProcessor,
+    Dinov2ForImageClassification,
+    Trainer,
+    TrainingArguments,
+)
+
+logger = logging.getLogger(__name__)
+
+MODEL = "facebook/dinov2-base-imagenet1k-1-layer"
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Train a DINOv2 model for font classification')
+    parser.add_argument('--data_dir', type=str, default=None,
+                      help='Directory containing the font dataset')
+    parser.add_argument('--output_dir', type=str, default=None,
+                      help='Directory to save the model')
+    parser.add_argument('--checkpoint', type=str, default=None,
+                      help='Path to checkpoint to resume training from')
+    parser.add_argument('--batch_size', type=int, default=32,
+                      help='Training and evaluation batch size')
+    parser.add_argument('--epochs', type=int, default=1,
+                      help='Number of training epochs')
+    parser.add_argument('--learning_rate', type=float, default=1e-4,
+                      help='Learning rate for training')
+    parser.add_argument('--lora_rank', type=int, default=8,
+                      help='LoRA rank for parameter-efficient fine-tuning')
+    parser.add_argument('--lora_alpha', type=int, default=16,
+                      help='LoRA alpha parameter')
+    parser.add_argument('--lora_dropout', type=float, default=0.1,
+                      help='LoRA dropout rate')
+    parser.add_argument('--test_size', type=float, default=0.1,
+                      help='Proportion of data to use for validation')
+    parser.add_argument('--seed', type=int, default=42,
+                      help='Random seed for reproducibility')
+    parser.add_argument('--log_level', type=str, default='INFO',
+                      choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
+                      help='Logging level')
+    parser.add_argument('--huggingface_model_name', type=str, default=None,
+                      help='Name of the model to push to the Hub')
+    return parser.parse_args()
+
+
+def load_checkpoint_with_size_mismatch_handling(base_model, checkpoint_path, peft_config):
+    """
+    Load PEFT checkpoint with automatic handling of size mismatches.
+    This uses PEFT's built-in loading but with strict=False to handle size mismatches gracefully.
+
+    Basically, if we have a different number of labels than in the checkpoint, we re-initialize the classifier head to relearn it.
+    
+    Args:
+        base_model: The base model with the new classifier size
+        checkpoint_path: Path to the checkpoint
+        peft_config: LoraConfig object with the desired configuration
+    
+    Returns:
+        PeftModel with loaded weights (mismatched layers will be skipped)
+    """
+    logger.info(f"Loading checkpoint with automatic size mismatch handling: {checkpoint_path}")
+    
+    try:
+        # Try the normal PEFT loading first
+        model = PeftModel.from_pretrained(
+            base_model,
+            checkpoint_path,
+            is_trainable=True
+        )
+        logger.info("Successfully loaded checkpoint without size mismatches")
+        return model
+    except Exception as e:
+        logger.info(f"Standard loading failed ({str(e)}), using fallback loading method")
+        
+        # Fallback: Create fresh PEFT model and load compatible weights
+        # Note: PeftModel.from_pretrained might have partially modified base_model before failing,
+        # so we recreate a clean base model to avoid double-loading warnings
+        fresh_base = Dinov2ForImageClassification.from_pretrained(
+            MODEL,
+            num_labels=base_model.config.num_labels,
+            ignore_mismatched_sizes=True,
+        )
+        
+        model = get_peft_model(fresh_base, peft_config)
+        
+        # Load checkpoint state dict
+        checkpoint_file = os.path.join(checkpoint_path, "adapter_model.safetensors")
+
+        if not os.path.exists(checkpoint_file):
+            raise ValueError(f"Checkpoint file {checkpoint_file} does not exist")
+        
+        checkpoint_state_dict = {}
+        with safe_open(checkpoint_file, framework="pt", device="cpu") as f:
+            for key in f.keys():
+                checkpoint_state_dict[key] = f.get_tensor(key)
+    
+        # Load only compatible weights
+        missing_keys, unexpected_keys = model.load_state_dict(checkpoint_state_dict, strict=False)
+        
+        logger.info(f"Loaded checkpoint with {len(missing_keys)} missing keys and {len(unexpected_keys)} unexpected keys")
+        logger.info(f"The following keys were in the checkpoint but are now missing: {missing_keys}")
+        logger.info(f"The following keys are new i.e. unexpected: {unexpected_keys}")
+        logger.info("Missing keys (likely new classifier parameters): will be randomly initialized")
+        
+        return model
+
+def get_inference_transform(processor: AutoImageProcessor, size: int):
+    """Get the raw validation transform for direct inference on PIL images."""
+    normalize   = T.Normalize(mean=processor.image_mean, std=processor.image_std)
+
+    to_rgb = T.Lambda(lambda img: img.convert('RGB'))
+ 
+    def pad_to_square(img):
+        w, h = img.size
+        max_size = max(w, h)
+        pad_w = (max_size - w) // 2
+        pad_h = (max_size - h) // 2
+        padding = (pad_w, pad_h, max_size - w - pad_w, max_size - h - pad_h)
+        return T.Pad(padding, fill=0)(img)
+
+    aug     = T.Compose([
+        to_rgb,
+        pad_to_square,
+        T.Resize(size),
+        T.ToTensor(), 
+        normalize
+    ])
+
+    return aug
+    
+
+
+def get_transform(processor: AutoImageProcessor, size: int):
+    aug = get_inference_transform(processor, size)
+
+    def transform(example, train=True):
+        # The dataset uses 'image' as the key for PIL images
+        # Use the processor directly - it handles pad_to_square + standard preprocessing
+        inputs = processor(images=example["image"], return_tensors="pt")
+        example["pixel_values"] = inputs["pixel_values"].squeeze(0)  # Remove batch dimension for dataset
+        return example
+
+    return transform
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    
+    # Configure logging with timestamps
+    logging.basicConfig(
+        level=args.log_level,
+        format='%(asctime)s - %(levelname)s - %(message)s - %(filename)s:%(lineno)d',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    )
+    
+    ######################################################################
+    # Directory layout expected by ImageFolder:
+    #   fonts/
+    #     ├─ Arial/
+    #     │   ├─ img001.png
+    #     │   └─ ...
+    #     ├─ TimesNewRoman/
+    #     └─ ...
+
+    logger.info(f"Loading dataset from {args.data_dir}")
+    # Get label names from directory names and sort them alphabetically
+    # to match the order used by the imagefolder dataset loader
+    label_names = sorted(os.listdir(f"{args.data_dir}/train"))
+    logger.info(f"Found {len(label_names)} labels")
+
+    if len(label_names) <= 1:
+        raise ValueError(f"Expected at least 2 labels, got {label_names=}, imagefolder will not label the dataset if there are less than 2 labels.")
+
+    # READ: the label ids assigned are in alphabetical order.
+    train_dataset = None
+    test_dataset = None
+    
+
+    logger.info("Setting up image processor and augmentations")
+    processor   = AutoImageProcessor.from_pretrained(MODEL)  # 224 px
+    size        = processor.size["shortest_edge"]            # 224 by default
+    
+    if args.epochs > 0:
+        dataset = load_dataset(
+            "imagefolder",
+            data_dir=args.data_dir,
+        )
+        
+        logger.info(f"Train size: {len(dataset['train'])}, Validation size: {len(dataset['test'])}")
+
+        transform = get_transform(processor, size)
+
+        logger.info("Applying data transformations")
+        train_dataset = dataset["train"].map(
+            lambda x: transform(x, train=True),
+            remove_columns=["image"],
+            desc="Transforming training data"
+        )
+        test_dataset = dataset["test"].map(
+            lambda x: transform(x, train=False),
+            remove_columns=["image"],
+            desc="Transforming test data"
+        )
+        
+        # Set the format to torch tensors
+        train_dataset.set_format(type="torch", columns=["pixel_values", "label"])
+        test_dataset.set_format(type="torch", columns=["pixel_values", "label"])
+        
+        logger.info("Data preprocessing complete")
+
+    logger.info("Loading DINOv2 model")
+    
+    base = Dinov2ForImageClassification.from_pretrained(
+            MODEL,
+            num_labels=len(label_names),
+            ignore_mismatched_sizes=True,
+        )
+
+    logger.info("Configuring LoRA adapters")
+    peft_cfg = LoraConfig(
+        r             = args.lora_rank,
+        lora_alpha    = args.lora_alpha,
+        target_modules = ["query", "value"],  # Q & V proj in ViT blocks
+        lora_dropout  = args.lora_dropout,
+        bias          = "none",
+        modules_to_save = ["classifier"],  # IMPORTANT: Save classification head too!
+    )
+
+    if args.checkpoint:
+        model = load_checkpoint_with_size_mismatch_handling(base, args.checkpoint, peft_cfg)
+    else:
+        model  = get_peft_model(base, peft_cfg)        # fresh LoRA wrap
+    
+    model.print_trainable_parameters()
+
+    def collate(batch):
+        # The transform function has already converted images to tensors and stored them in pixel_values
+        pixel_values = torch.stack([item["pixel_values"] for item in batch])
+        labels = torch.tensor([item["label"] for item in batch])
+        return {"pixel_values": pixel_values, "labels": labels}
+
+    # Add compute_metrics function for accuracy calculation
+    def compute_metrics(eval_pred):
+        predictions, labels = eval_pred
+        predictions = predictions.argmax(axis=-1)
+        accuracy = (predictions == labels).mean()
+        return {"accuracy": accuracy}
+
+    logger.info("Setting up training arguments")
+    # Check if we're on MPS (Apple Silicon)
+    device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
+    logger.info(f"Using device: {device}")
+    
+    training_args = TrainingArguments(
+        output_dir          = args.output_dir,
+        per_device_train_batch_size = args.batch_size,
+        per_device_eval_batch_size  = args.batch_size,
+        # Tell Trainer which key in each batch holds the ground‑truth labels.
+        # Without it (especially with PEFT/LoRA wrappers), Trainer thinks there
+        # are no labels, skips compute_metrics, and never logs eval_accuracy.
+        label_names=["labels"],
+        eval_strategy      = "steps" if args.epochs > 0 else "no",
+        eval_steps         = 500,
+        save_strategy      = "steps" if args.epochs > 0 else "no",
+        save_steps         = 500,
+        num_train_epochs   = args.epochs,
+        learning_rate      = args.learning_rate,
+        weight_decay       = 0.05,
+        fp16               = device.type == "cuda",
+        save_total_limit   = 3,
+        logging_dir        = os.path.join(args.output_dir, "logs") if args.output_dir else None,
+        logging_steps      = 10,
+        report_to          = "tensorboard",
+        load_best_model_at_end = True,
+        metric_for_best_model = "eval_accuracy",
+        greater_is_better = True,
+        # Pass the actual checkpoint path for proper resumption
+        resume_from_checkpoint = args.checkpoint if args.checkpoint else None,
+    )
+
+    trainer = Trainer(
+        model           = model,
+        args            = training_args,
+        train_dataset   = train_dataset,
+        eval_dataset    = test_dataset,
+        data_collator   = collate,
+        compute_metrics = compute_metrics,
+    )
+
+    logger.info("Starting training")
+    if args.checkpoint:
+        logger.info(f"Resuming training from checkpoint: {args.checkpoint}")
+    
+    if args.epochs > 0:
+        trainer.train()
+    logger.info("Training complete")
+    
+    # Saves the result model to the output directory
+    # The reason this is important is if we configure load_best_model_at_end=True,
+    # the best model will be saved out of all checkpoints.
+    # So, even though the trainer already saves the last model as a checkpoint, that one is not necessarily the best.
+    if args.output_dir:
+        logger.info("Saving result model to the output directory")
+        trainer.save_model(f"{args.output_dir}/result_model")
+
+    if args.huggingface_model_name:
+        logger.info(f"Pushing model to the Hub: {args.huggingface_model_name}")
+
+        trainer.hub_model_id = args.huggingface_model_name
+
+        with tempfile.TemporaryDirectory() as tmp:
+            # Merge the PEFT weights into the base model so that we upload an independent complete model.
+            merged = trainer.model.merge_and_unload()
+            id2label = {i: name for i, name in enumerate(label_names)}
+            label2id = {name: i for i, name in enumerate(label_names)}
+
+            merged.config.id2label = id2label
+            merged.config.label2id = label2id
+            merged.config.pipeline_tag = "image-classification"
+            merged.save_pretrained(tmp, safe_serialization=True)
+            processor.save_pretrained(tmp)
+
+            # bundle handler and code
+            shutil.copy("train_model.py", tmp)
+            shutil.copy("handler.py", tmp)
+            Path(tmp, "requirements.txt").write_text("torchvision>=0.19\nPillow>=10\n")
+
+            HfApi().upload_folder(
+                repo_id=args.huggingface_model_name, 
+                folder_path=tmp,
+                commit_message="Add merged model + processor",
+                token=os.environ["HUGGINGFACE_API_KEY"],
+            )