models/base_loader.py

"""
models/base_loader.py
─────────────────────
Abstract contract that every stage loader must satisfy.
Stage runners in pipeline.py interact only with this interface —
swapping models is just swapping loaders.
"""

from __future__ import annotations

import abc
import logging
from typing import Any

import torch

logger = logging.getLogger(__name__)


class BaseLoader(abc.ABC):
    """
    Lifecycle
    ---------
    1. Instantiate with model_id and kwargs.
    2. Call  .load()  once to download/initialise weights.
    3. Call  .run(**inputs)  as many times as needed.
    4. Call  .unload()  when done to free memory.
    """

    def __init__(self, model_id: str, device: torch.device, **kwargs: Any) -> None:
        self.model_id = model_id
        self.device = device
        self.kwargs = kwargs
        self._loaded = False
        self.logger = logging.getLogger(self.__class__.__name__)

    # ── Required interface ────────────────────────────────────────────────────

    @abc.abstractmethod
    def load(self) -> None:
        """Download weights and move model to self.device."""

    @abc.abstractmethod
    def run(self, **inputs: Any) -> dict[str, Any]:
        """
        Execute the model.

        Inputs and outputs are stage-specific dictionaries; see each
        concrete loader for the expected keys.
        """

    # ── Optional hooks ────────────────────────────────────────────────────────

    def unload(self) -> None:
        """Release model weights and clear GPU cache. Override if needed."""
        from utils.device import clear_cache
        for attr in ("model", "pipe", "processor", "feature_extractor"):
            if hasattr(self, attr):
                delattr(self, attr)
        self._loaded = False
        clear_cache()
        self.logger.info("Unloaded %s", self.__class__.__name__)

    def is_loaded(self) -> bool:
        return self._loaded

    # ── Helpers ───────────────────────────────────────────────────────────────

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(model_id={self.model_id!r}, device={self.device})"