Spaces:
Sleeping
Sleeping
| """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)" | |
| 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).", | |
| ), | |
| ] | |
| 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() | |
| 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 | |
| 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}") | |
| def _null_ctx(): | |
| yield | |