Upload folder using huggingface_hub

README.md

@@ -71,11 +71,47 @@ git push -u origin main
 The Space will build on a free CPU runtime by default. For faster inference you can
 upgrade the Space hardware to a small GPU (`T4`, `A10G`, etc.) in the Space settings.
 
+## Training a new LoRA adapter
+
+The demo also supports LoRA adapters that add new tasks on top of the frozen
+backbone. Train one with:
+
+```bash
+python train_lora.py \
+  --rank 8 --alpha 16 --target-modules query value \
+  --epochs 5 --batch-size 64 --lr 5e-4 \
+  --push-to-hub <your-username>/vit-tiny-lora-food101
+```
+
+The script freezes the base weights, injects a low-rank \(\Delta W\) into the
+attention projections, and trains a new classification head. Because the
+original weights are untouched, disabling the adapter at inference time
+recovers the original ImageNet-1k model exactly.
+
+Useful flags: `--max-train-samples N` (quick smoke test), `--eval-only`
+(metrics-only pass), `--dataset-id` (any HF image classification dataset with
+`image` / `label` features).
+
+## Adapters loaded at runtime
+
+`adapters.py` holds the registry of adapters the Gradio app pulls in at
+startup. Add more entries to expose additional tasks in the UI:
+
+| Name      | Hub repo                                 | Dataset   | Classes |
+|-----------|------------------------------------------|-----------|---------|
+| food101   | `turhancan97/vit-tiny-lora-food101`      | Food-101  | 101     |
+
+Adapters that fail to load (e.g. repo not yet pushed) are logged and skipped;
+the app still starts with whatever is reachable. The UI gains a "Compare: Base
+vs LoRA" tab whenever at least one adapter is loaded.
+
 ## Project layout
 
 ```
 .
 ├── app.py              # Gradio Blocks app
+├── adapters.py         # LoRA adapter registry
+├── train_lora.py       # LoRA fine-tuning CLI
 ├── requirements.txt    # Python deps
 ├── examples/           # Sample images used in the UI
 │   ├── bird.jpg

adapters.py

@@ -0,0 +1,209 @@
+"""Registry of LoRA adapters for the Gradio demo.
+
+Loads the pretrained base model once and wraps it with a `PeftModel` that
+holds one or more named LoRA adapters. Each adapter also carries its own
+classification head (saved separately as ``classifier.pt`` at training time),
+so the demo can swap *both* the adapter and the head per request.
+
+Inference modes:
+    * Base (ImageNet-1k): ``model.disable_adapter()`` + original head.
+    * Adapter N:          ``model.set_adapter(name)`` + the adapter's head.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Iterable
+
+import torch
+from torch import nn
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+
+BASE_MODEL_ID = "WinKawaks/vit-tiny-patch16-224"
+BASE_TASK_NAME = "imagenet"
+BASE_TASK_LABEL = "ImageNet-1k (base)"
+
+
+@dataclass
+class AdapterSpec:
+    name: str
+    repo_id: str
+    display_name: str
+    description: str = ""
+    id2label: dict[int, str] = field(default_factory=dict)
+
+
+ADAPTERS: list[AdapterSpec] = [
+    AdapterSpec(
+        name="food101",
+        repo_id="turhancan97/vit-tiny-lora-food101",
+        display_name="Food-101 (LoRA)",
+        description="Fine-tuned LoRA adapter on the Food-101 dataset (101 classes).",
+    ),
+]
+
+
+@dataclass
+class TaskInfo:
+    key: str
+    display_name: str
+    id2label: dict[int, str]
+    is_base: bool
+
+
+def _fetch_file(repo_id: str, filename: str) -> Path:
+    local = Path(repo_id)
+    if local.exists() and (local / filename).exists():
+        return local / filename
+    from huggingface_hub import hf_hub_download
+
+    return Path(hf_hub_download(repo_id=repo_id, filename=filename))
+
+
+class AdapterRegistry:
+    """Loads the base model plus any registered adapters lazily and on demand."""
+
+    def __init__(
+        self,
+        adapters: Iterable[AdapterSpec] | None = None,
+        base_model_id: str = BASE_MODEL_ID,
+        device: str | None = None,
+    ) -> None:
+        if adapters is None:
+            adapters = ADAPTERS
+        self.base_model_id = base_model_id
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.processor = AutoImageProcessor.from_pretrained(base_model_id, use_fast=True)
+        self.base_model = AutoModelForImageClassification.from_pretrained(base_model_id)
+        self.base_id2label: dict[int, str] = dict(self.base_model.config.id2label)
+
+        self._original_classifier = copy.deepcopy(self.base_model.classifier)
+
+        self._model = None
+        self._loaded: dict[str, AdapterSpec] = {}
+        self._classifiers: dict[str, nn.Module] = {}
+        self._current_adapter: str | None = None
+
+        for spec in adapters:
+            self._try_load_adapter(spec)
+
+        if self._model is None:
+            self._model = self.base_model
+        self._model.to(self.device).eval()
+        self._original_classifier.to(self.device).eval()
+        for head in self._classifiers.values():
+            head.to(self.device).eval()
+
+    @property
+    def model(self):
+        return self._model
+
+    def available_tasks(self) -> list[TaskInfo]:
+        tasks = [TaskInfo(
+            key=BASE_TASK_NAME,
+            display_name=BASE_TASK_LABEL,
+            id2label=self.base_id2label,
+            is_base=True,
+        )]
+        for spec in self._loaded.values():
+            tasks.append(TaskInfo(
+                key=spec.name,
+                display_name=spec.display_name,
+                id2label=spec.id2label,
+                is_base=False,
+            ))
+        return tasks
+
+    def adapter_task_keys(self) -> list[str]:
+        return list(self._loaded.keys())
+
+    def get_id2label(self, task_key: str) -> dict[int, str]:
+        if task_key == BASE_TASK_NAME:
+            return self.base_id2label
+        return self._loaded[task_key].id2label
+
+    def _classifier_host(self) -> nn.Module:
+        """Return the module that actually *owns* the `classifier` submodule.
+
+        `PeftModel` exposes `.classifier` via attribute proxying on top of
+        `base_model.model.classifier`, but only the latter is what gets called
+        during forward. Traverse until we find a module where `classifier` is
+        a registered child (present in `_modules`).
+        """
+        visited: set[int] = set()
+        m = self._model
+        while id(m) not in visited:
+            visited.add(id(m))
+            if "classifier" in getattr(m, "_modules", {}):
+                return m
+            inner = getattr(m, "model", None) or getattr(m, "base_model", None)
+            if inner is None or inner is m:
+                break
+            m = inner
+        return self.base_model
+
+    def _active_classifier(self) -> nn.Module:
+        return self._classifier_host().classifier
+
+    def _set_classifier(self, head: nn.Module) -> None:
+        self._classifier_host().classifier = head
+
+    @contextmanager
+    def use_task(self, task_key: str):
+        """Activate the right (adapter, head) pair for a single forward pass."""
+        prev_head = self._active_classifier()
+        if task_key == BASE_TASK_NAME:
+            self._set_classifier(self._original_classifier)
+            ctx = self._model.disable_adapter() if hasattr(self._model, "disable_adapter") else _null_ctx()
+        else:
+            spec = self._loaded[task_key]
+            self._set_classifier(self._classifiers[spec.name])
+            if hasattr(self._model, "set_adapter") and self._current_adapter != task_key:
+                self._model.set_adapter(task_key)
+                self._current_adapter = task_key
+            ctx = _null_ctx()
+
+        try:
+            with ctx:
+                yield
+        finally:
+            self._set_classifier(prev_head)
+
+    def _try_load_adapter(self, spec: AdapterSpec) -> None:
+        try:
+            labels_path = _fetch_file(spec.repo_id, "labels.json")
+            spec.id2label = {int(k): v for k, v in json.loads(labels_path.read_text()).items()}
+
+            classifier_path = _fetch_file(spec.repo_id, "classifier.pt")
+            head_state = torch.load(classifier_path, map_location="cpu", weights_only=True)
+            hidden_size = self.base_model.config.hidden_size
+            head = nn.Linear(hidden_size, len(spec.id2label))
+            head.load_state_dict(head_state)
+
+            from peft import PeftModel
+
+            if self._model is None or not hasattr(self._model, "load_adapter"):
+                self._model = PeftModel.from_pretrained(
+                    self.base_model, spec.repo_id, adapter_name=spec.name,
+                )
+            else:
+                self._model.load_adapter(spec.repo_id, adapter_name=spec.name)
+
+            self._classifiers[spec.name] = head
+            self._loaded[spec.name] = spec
+            print(f"[adapters] loaded '{spec.name}' from {spec.repo_id} "
+                  f"({len(spec.id2label)} classes)")
+        except Exception as exc:
+            print(f"[adapters] could not load '{spec.name}' from {spec.repo_id}: "
+                  f"{type(exc).__name__}: {exc}")
+
+
+@contextmanager
+def _null_ctx():
+    yield

app.py

@@ -1,4 +1,12 @@
-"""Gradio demo: Tiny ViT (ImageNet-1k) image classification with top-k + threshold."""
+"""Gradio demo: ViT image classification with pluggable LoRA adapters.
+
+Tabs:
+    Classify          - pick a task (ImageNet-1k base or any loaded LoRA adapter)
+                        and get top-k predictions with a confidence threshold.
+    Compare: Base vs  - side-by-side predictions from the same backbone with
+    LoRA                the adapter disabled vs enabled. Makes it obvious what
+                        LoRA added without touching the pretrained weights.
+"""
 
 from __future__ import annotations
 
@@ -28,31 +36,40 @@ _gc_utils._json_schema_to_python_type = _safe_json_schema_to_python_type
 
 import gradio as gr
 import torch
-from PIL import Image
-from transformers import AutoImageProcessor, AutoModelForImageClassification
 
-MODEL_ID = "WinKawaks/vit-tiny-patch16-224"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+from adapters import BASE_TASK_NAME, AdapterRegistry
+
+
+registry = AdapterRegistry()
+model = registry.model
+processor = registry.processor
+DEVICE = registry.device
 
-processor = AutoImageProcessor.from_pretrained(MODEL_ID)
-model = AutoModelForImageClassification.from_pretrained(MODEL_ID).to(DEVICE).eval()
-id2label = model.config.id2label
+TASKS = registry.available_tasks()
+TASK_CHOICES = [(t.display_name, t.key) for t in TASKS]
+ADAPTER_TASKS = [t for t in TASKS if not t.is_base]
+ADAPTER_CHOICES = [(t.display_name, t.key) for t in ADAPTER_TASKS]
 
 EXAMPLES_DIR = Path(__file__).parent / "examples"
 EXAMPLE_IMAGES = sorted(str(p) for p in EXAMPLES_DIR.glob("*.jpg")) if EXAMPLES_DIR.exists() else []
 
 
-@torch.inference_mode()
-def classify(image, top_k, threshold):
-    if image is None:
-        return {}
+def _to_inputs(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
+    return processor(images=image, return_tensors="pt").to(DEVICE)
+
 
-    inputs = processor(images=image, return_tensors="pt").to(DEVICE)
-    logits = model(**inputs).logits[0]
+@torch.inference_mode()
+def classify(image, task_key, top_k, threshold):
+    if image is None:
+        return {}
+    inputs = _to_inputs(image)
+    with registry.use_task(task_key):
+        logits = model(**inputs).logits[0]
     probs = torch.softmax(logits, dim=-1)
 
+    id2label = registry.get_id2label(task_key)
     k = max(1, min(int(top_k), probs.shape[-1]))
     top = torch.topk(probs, k=k)
     return {
@@ -62,60 +79,140 @@ def classify(image, top_k, threshold):
     }
 
 
-with gr.Blocks(title="Tiny ViT — ImageNet-1k classifier", theme=gr.themes.Soft()) as demo:
+def classify_compare(image, adapter_key, top_k, threshold):
+    if image is None or not adapter_key:
+        return {}, {}
+    base = classify(image, BASE_TASK_NAME, top_k, threshold)
+    adapter = classify(image, adapter_key, top_k, threshold)
+    return base, adapter
+
+
+default_task = ADAPTER_CHOICES[0][1] if ADAPTER_CHOICES else BASE_TASK_NAME
+adapter_names_loaded = [a.display_name for a in ADAPTER_TASKS]
+
+with gr.Blocks(title="ViT + LoRA image classifier", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         f"""
-        # Tiny ViT — ImageNet-1k classifier
+        # ViT image classifier with swappable LoRA adapters
 
-        Upload an image and get top-k predictions from
-        [`{MODEL_ID}`](https://huggingface.co/{MODEL_ID}), a tiny Vision Transformer
-        (~5.7M params) pretrained on ImageNet-1k.
+        Backbone: [`{registry.base_model_id}`](https://huggingface.co/{registry.base_model_id}) (ViT-tiny, ~5.7M params).
+        Device: **{DEVICE.upper()}**.
+        Loaded adapters: **{', '.join(adapter_names_loaded) if adapter_names_loaded else 'none (base model only)'}**.
 
-        Running on **{DEVICE.upper()}**.
+        Pick a task on the left to run top-k classification. The base weights
+        are never mutated: the adapter adds a low-rank \(\\Delta W\) and a new
+        classification head, toggled on/off per request.
         """
     )
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            image = gr.Image(
-                type="pil",
-                label="Input image",
-                sources=["upload", "clipboard", "webcam"],
-                height=360,
-            )
+    with gr.Tabs():
+        with gr.Tab("Classify"):
             with gr.Row():
-                top_k = gr.Slider(
-                    minimum=1, maximum=10, value=5, step=1,
-                    label="Top-k", info="Number of predictions to show",
+                with gr.Column(scale=1):
+                    image = gr.Image(
+                        type="pil", label="Input image",
+                        sources=["upload", "clipboard", "webcam"], height=360,
+                    )
+                    task = gr.Radio(
+                        choices=TASK_CHOICES, value=BASE_TASK_NAME,
+                        label="Task", info="Switch between the base model and any loaded LoRA adapter.",
+                    )
+                    with gr.Row():
+                        top_k = gr.Slider(
+                            minimum=1, maximum=10, value=5, step=1,
+                            label="Top-k", info="Number of predictions to show",
+                        )
+                        threshold = gr.Slider(
+                            minimum=0.0, maximum=1.0, value=0.0, step=0.01,
+                            label="Confidence threshold",
+                            info="Hide predictions below this probability",
+                        )
+                    submit = gr.Button("Classify", variant="primary")
+
+                with gr.Column(scale=1):
+                    output = gr.Label(num_top_classes=10, label="Predictions")
+
+            if EXAMPLE_IMAGES:
+                gr.Examples(
+                    examples=[[p, BASE_TASK_NAME, 5, 0.0] for p in EXAMPLE_IMAGES],
+                    inputs=[image, task, top_k, threshold],
+                    outputs=output,
+                    fn=classify,
+                    cache_examples=False,
+                    label="Example images",
                 )
-                threshold = gr.Slider(
-                    minimum=0.0, maximum=1.0, value=0.0, step=0.01,
-                    label="Confidence threshold",
-                    info="Hide predictions below this probability",
-                )
-            submit = gr.Button("Classify", variant="primary")
-
-        with gr.Column(scale=1):
-            output = gr.Label(num_top_classes=10, label="Predictions")
 
-    if EXAMPLE_IMAGES:
-        gr.Examples(
-            examples=[[p, 5, 0.0] for p in EXAMPLE_IMAGES],
-            inputs=[image, top_k, threshold],
-            outputs=output,
-            fn=classify,
-            cache_examples=False,
-            label="Example images",
-        )
-
-    submit.click(fn=classify, inputs=[image, top_k, threshold], outputs=output)
-    image.change(fn=classify, inputs=[image, top_k, threshold], outputs=output)
+            classify_args = dict(fn=classify, inputs=[image, task, top_k, threshold], outputs=output)
+            submit.click(**classify_args)
+            image.change(**classify_args)
+            task.change(**classify_args)
+            top_k.change(**classify_args)
+            threshold.change(**classify_args)
+
+        if ADAPTER_TASKS:
+            with gr.Tab("Compare: Base vs LoRA"):
+                gr.Markdown(
+                    "Runs the **same image** through the shared backbone twice: "
+                    "once with the LoRA adapter disabled (original ImageNet-1k head) "
+                    "and once with the selected adapter enabled."
+                )
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        image2 = gr.Image(
+                            type="pil", label="Input image",
+                            sources=["upload", "clipboard", "webcam"], height=360,
+                        )
+                        adapter_sel = gr.Dropdown(
+                            choices=ADAPTER_CHOICES, value=default_task,
+                            label="Adapter", interactive=True,
+                        )
+                        with gr.Row():
+                            top_k2 = gr.Slider(
+                                minimum=1, maximum=10, value=5, step=1, label="Top-k",
+                            )
+                            threshold2 = gr.Slider(
+                                minimum=0.0, maximum=1.0, value=0.0, step=0.01,
+                                label="Confidence threshold",
+                            )
+                        submit2 = gr.Button("Compare", variant="primary")
+
+                with gr.Row():
+                    base_out = gr.Label(num_top_classes=10, label="Base (ImageNet-1k)")
+                    adapter_out = gr.Label(num_top_classes=10, label="Adapter prediction")
+
+                if EXAMPLE_IMAGES:
+                    gr.Examples(
+                        examples=[[p, default_task, 5, 0.0] for p in EXAMPLE_IMAGES],
+                        inputs=[image2, adapter_sel, top_k2, threshold2],
+                        outputs=[base_out, adapter_out],
+                        fn=classify_compare,
+                        cache_examples=False,
+                        label="Example images",
+                    )
+
+                compare_args = dict(
+                    fn=classify_compare,
+                    inputs=[image2, adapter_sel, top_k2, threshold2],
+                    outputs=[base_out, adapter_out],
+                )
+                submit2.click(**compare_args)
+                image2.change(**compare_args)
+                adapter_sel.change(**compare_args)
+                top_k2.change(**compare_args)
+                threshold2.change(**compare_args)
+        else:
+            with gr.Tab("Compare: Base vs LoRA"):
+                gr.Markdown(
+                    "_No LoRA adapters are currently loaded._\n\n"
+                    "Train one with `python train_lora.py --push-to-hub <user>/<repo>` "
+                    "and register it in `adapters.py`."
+                )
 
 
 if __name__ == "__main__":
     demo.launch(
-        # server_name="0.0.0.0",
-        # server_port=7860,
-        # show_api=False,
-        # share=False,
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_api=False,
+        share=False,
     )

requirements.txt

@@ -2,3 +2,7 @@ gradio>=5.0.0
 transformers>=4.44.0
 torch>=2.2.0
 Pillow>=10.0.0
+peft>=0.13.0
+datasets>=2.18.0
+accelerate>=0.30.0
+scikit-learn

train_lora.py

@@ -0,0 +1,275 @@
+"""Fine-tune a LoRA adapter on Food-101 for a ViT image classifier.
+
+Preserves the original pretrained weights (LoRA is additive) and saves the
+adapter + the new classification head as a single PEFT-format artifact.
+
+Example:
+
+    python train_lora.py \\
+        --rank 8 --alpha 16 --target-modules query value \\
+        --epochs 5 --batch-size 64 --lr 5e-4 \\
+        --push-to-hub turhancan97/vit-tiny-lora-food101
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model
+from PIL import Image
+from torchvision import transforms
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+    Trainer,
+    TrainingArguments,
+)
+
+
+@dataclass
+class Args:
+    model_id: str
+    dataset_id: str
+    output_dir: str
+    rank: int
+    alpha: int
+    dropout: float
+    target_modules: list[str]
+    lr: float
+    batch_size: int
+    eval_batch_size: int
+    epochs: int
+    warmup_ratio: float
+    weight_decay: float
+    seed: int
+    push_to_hub: str | None
+    max_train_samples: int | None
+    max_eval_samples: int | None
+    eval_only: bool
+    num_workers: int
+
+
+def parse_args() -> Args:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model-id", default="WinKawaks/vit-tiny-patch16-224")
+    p.add_argument("--dataset-id", default="food101")
+    p.add_argument("--output-dir", default="adapters/vit-tiny-lora-food101")
+    p.add_argument("--rank", type=int, default=8)
+    p.add_argument("--alpha", type=int, default=16)
+    p.add_argument("--dropout", type=float, default=0.1)
+    p.add_argument(
+        "--target-modules", nargs="+", default=["query", "value"],
+        help="Substring patterns matched against module names for LoRA injection.",
+    )
+    p.add_argument("--lr", type=float, default=5e-4)
+    p.add_argument("--batch-size", type=int, default=64)
+    p.add_argument("--eval-batch-size", type=int, default=128)
+    p.add_argument("--epochs", type=int, default=1)
+    p.add_argument("--warmup-ratio", type=float, default=0.03)
+    p.add_argument("--weight-decay", type=float, default=0.0)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--push-to-hub", default=None, help="e.g. 'user/vit-tiny-lora-food101'")
+    p.add_argument("--max-train-samples", type=int, default=None, help="Smoke-test subset size.")
+    p.add_argument("--max-eval-samples", type=int, default=None)
+    p.add_argument("--eval-only", action="store_true")
+    p.add_argument("--num-workers", type=int, default=4)
+    ns = p.parse_args()
+    return Args(**{k.replace("-", "_"): v for k, v in vars(ns).items()})
+
+
+def build_transforms(processor: AutoImageProcessor):
+    size = processor.size.get("height") or processor.size.get("shortest_edge") or 224
+    mean = processor.image_mean
+    std = processor.image_std
+
+    train_tf = transforms.Compose([
+        transforms.RandomResizedCrop(size, scale=(0.8, 1.0)),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=mean, std=std),
+    ])
+    eval_tf = transforms.Compose([
+        transforms.Resize(int(size * 256 / 224)),
+        transforms.CenterCrop(size),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=mean, std=std),
+    ])
+    return train_tf, eval_tf
+
+
+def _ensure_rgb(img):
+    if isinstance(img, Image.Image):
+        return img.convert("RGB") if img.mode != "RGB" else img
+    return Image.fromarray(np.asarray(img)).convert("RGB")
+
+
+def make_transform_fn(tf):
+    def _apply(batch):
+        batch["pixel_values"] = [tf(_ensure_rgb(img)) for img in batch["image"]]
+        return batch
+    return _apply
+
+
+def collate_fn(examples):
+    pixel_values = torch.stack([ex["pixel_values"] for ex in examples])
+    labels = torch.tensor([ex["label"] for ex in examples], dtype=torch.long)
+    return {"pixel_values": pixel_values, "labels": labels}
+
+
+def compute_metrics_topk(eval_pred):
+    logits, labels = eval_pred
+    logits = torch.as_tensor(logits)
+    labels = torch.as_tensor(labels)
+    top1 = (logits.argmax(dim=-1) == labels).float().mean().item()
+    k = min(5, logits.shape[-1])
+    topk = logits.topk(k=k, dim=-1).indices
+    top5 = (topk == labels.unsqueeze(-1)).any(dim=-1).float().mean().item()
+    return {"top1_accuracy": top1, "top5_accuracy": top5}
+
+
+def main():
+    args = parse_args()
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"[1/5] Loading dataset: {args.dataset_id}")
+    ds = load_dataset(args.dataset_id)
+    train_split = "train" if "train" in ds else list(ds.keys())[0]
+    eval_split = "validation" if "validation" in ds else ("test" if "test" in ds else train_split)
+    train_ds = ds[train_split]
+    eval_ds = ds[eval_split]
+
+    label_feature = train_ds.features["label"]
+    num_labels = label_feature.num_classes
+    id2label = {i: label_feature.int2str(i) for i in range(num_labels)}
+    label2id = {v: k for k, v in id2label.items()}
+    print(f"    train={len(train_ds)}  eval={len(eval_ds)}  num_labels={num_labels}")
+
+    if args.max_train_samples:
+        train_ds = train_ds.shuffle(seed=args.seed).select(range(min(args.max_train_samples, len(train_ds))))
+    if args.max_eval_samples:
+        eval_ds = eval_ds.shuffle(seed=args.seed).select(range(min(args.max_eval_samples, len(eval_ds))))
+
+    print(f"[2/5] Loading base model: {args.model_id}")
+    processor = AutoImageProcessor.from_pretrained(args.model_id, use_fast=True)
+    base_model = AutoModelForImageClassification.from_pretrained(
+        args.model_id,
+        num_labels=num_labels,
+        id2label=id2label,
+        label2id=label2id,
+        ignore_mismatched_sizes=True,
+    )
+
+    train_tf, eval_tf = build_transforms(processor)
+    train_ds.set_transform(make_transform_fn(train_tf))
+    eval_ds.set_transform(make_transform_fn(eval_tf))
+
+    print(f"[3/5] Wrapping with LoRA: rank={args.rank}, alpha={args.alpha}, "
+          f"target_modules={args.target_modules}")
+    lora_cfg = LoraConfig(
+        r=args.rank,
+        lora_alpha=args.alpha,
+        lora_dropout=args.dropout,
+        target_modules=list(args.target_modules),
+        bias="none",
+    )
+    model = get_peft_model(base_model, lora_cfg)
+    # PEFT freezes every non-LoRA parameter by default. Unfreeze the classifier
+    # so the new task head can be trained. We save it separately after training
+    # (rather than via `modules_to_save`) so the adapter artifact stays portable
+    # across base models with different original head sizes.
+    classifier = model.base_model.model.classifier
+    for p in classifier.parameters():
+        p.requires_grad_(True)
+    trainable, total = model.get_nb_trainable_parameters()
+    print(f"    trainable params: {trainable:,} / {total:,}  ({100 * trainable / total:.2f}%)")
+
+    training_args = TrainingArguments(
+        output_dir=str(output_dir / "trainer"),
+        per_device_train_batch_size=args.batch_size,
+        per_device_eval_batch_size=args.eval_batch_size,
+        learning_rate=args.lr,
+        num_train_epochs=args.epochs,
+        warmup_ratio=args.warmup_ratio,
+        weight_decay=args.weight_decay,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        save_total_limit=1,
+        load_best_model_at_end=True,
+        metric_for_best_model="top1_accuracy",
+        greater_is_better=True,
+        logging_strategy="steps",
+        logging_steps=25,
+        fp16=torch.cuda.is_available(),
+        dataloader_num_workers=args.num_workers,
+        remove_unused_columns=False,
+        report_to="none",
+        seed=args.seed,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        data_collator=collate_fn,
+        compute_metrics=compute_metrics_topk,
+    )
+
+    if not args.eval_only:
+        print("[4/5] Training")
+        trainer.train()
+    else:
+        print("[4/5] Skipping training (--eval-only)")
+
+    print("[5/5] Evaluating on held-out split")
+    metrics = trainer.evaluate()
+    metrics["eval_samples"] = len(eval_ds)
+    print(json.dumps(metrics, indent=2))
+    (output_dir / "eval_metrics.json").write_text(json.dumps(metrics, indent=2))
+
+    print(f"Saving adapter to {output_dir}")
+    model.save_pretrained(str(output_dir))
+    processor.save_pretrained(str(output_dir))
+    (output_dir / "train_args.json").write_text(json.dumps(asdict(args), indent=2))
+    (output_dir / "labels.json").write_text(
+        json.dumps({str(i): id2label[i] for i in range(num_labels)}, indent=2)
+    )
+    torch.save(
+        {k: v.detach().cpu() for k, v in classifier.state_dict().items()},
+        output_dir / "classifier.pt",
+    )
+
+    if args.push_to_hub:
+        print(f"Pushing to Hugging Face Hub: {args.push_to_hub}")
+        model.push_to_hub(args.push_to_hub)
+        processor.push_to_hub(args.push_to_hub)
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi()
+            for extra in ["labels.json", "classifier.pt"]:
+                api.upload_file(
+                    path_or_fileobj=str(output_dir / extra),
+                    path_in_repo=extra,
+                    repo_id=args.push_to_hub,
+                    repo_type="model",
+                    commit_message=f"add {extra}",
+                )
+        except Exception as exc:
+            print(f"Warning: could not upload side-car files: {exc}")
+
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()