models/base.py

"""Common lifecycle contract for otherwise isolated model adapters."""

from __future__ import annotations

from abc import ABC, abstractmethod

from config import Settings
from core.runtime import cleanup_memory


class ModelAdapter(ABC):
    """A lazy model with explicit device release and full unload hooks."""

    name: str

    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self.device = "cpu"
        self.loaded = False

    def ensure_loaded(self, device: str) -> None:
        """Load once, or reactivate an already resident model."""
        if not self.loaded:
            self.load(device)
            self.loaded = True
        else:
            self.activate(device)
        self.device = device

    @abstractmethod
    def load(self, device: str) -> None:
        """Load model weights and processors."""

    def activate(self, device: str) -> None:
        """Move a cached model back to an accelerator when necessary."""

    @abstractmethod
    def release_gpu(self) -> None:
        """Move allocations off GPU after every inference call."""

    @abstractmethod
    def unload(self) -> None:
        """Release all model references before another model is loaded."""

    def mark_unloaded(self) -> None:
        self.loaded = False
        self.device = "cpu"
        cleanup_memory()