invoke-ai--invokeai/invokeai/backend/model_manager/load/load_base.py

# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
"""
Base class for model loading in InvokeAI.
"""

from abc import ABC, abstractmethod
from contextlib import contextmanager
from logging import Logger
from pathlib import Path
from typing import Any, Dict, Generator, Optional, Tuple

import torch

from invokeai.app.services.config import InvokeAIAppConfig
from invokeai.backend.model_manager.configs.factory import AnyModelConfig
from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord
from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_with_partial_load import (
    CachedModelWithPartialLoad,
)
from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType


class LoadedModelWithoutConfig:
    """Context manager object that mediates transfer from RAM<->VRAM.

    This is a context manager object that has two distinct APIs:

    1. Older API (deprecated):
    Use the LoadedModel object directly as a context manager.  It will move the model into VRAM (on CUDA devices), and
    return the model in a form suitable for passing to torch.
    Example:
    ```
    loaded_model_= loader.get_model_by_key('f13dd932', SubModelType('vae'))
    with loaded_model as vae:
      image = vae.decode(latents)[0]
    ```

    2. Newer API (recommended):
    Call the LoadedModel's `model_on_device()` method in a context. It returns a tuple consisting of a copy of the
    model's state dict in CPU RAM followed by a copy of the model in VRAM. The state dict is provided to allow LoRAs and
    other model patchers to return the model to its unpatched state without expensive copy and restore operations.

    Example:
    ```
    loaded_model_= loader.get_model_by_key('f13dd932', SubModelType('vae'))
    with loaded_model.model_on_device() as (state_dict, vae):
        image = vae.decode(latents)[0]
    ```

    The state_dict should be treated as a read-only object and never modified. Also be aware that some loadable models
    do not have a state_dict, in which case this value will be None.
    """

    def __init__(self, cache_record: CacheRecord, cache: ModelCache):
        self._cache_record = cache_record
        self._cache = cache

    def __enter__(self) -> AnyModel:
        self._cache.lock(self._cache_record, None)
        try:
            self.repair_required_tensors_on_device()
            return self.model
        except Exception:
            self._cache.unlock(self._cache_record)
            raise

    def __exit__(self, *args: Any, **kwargs: Any) -> None:
        self._cache.unlock(self._cache_record)

    @contextmanager
    def model_on_device(
        self, working_mem_bytes: Optional[int] = None
    ) -> Generator[Tuple[Optional[Dict[str, torch.Tensor]], AnyModel], None, None]:
        """Return a tuple consisting of the model's state dict (if it exists) and the locked model on execution device.

        :param working_mem_bytes: The amount of working memory to keep available on the compute device when loading the
            model.
        """
        self._cache.lock(self._cache_record, working_mem_bytes)
        try:
            self.repair_required_tensors_on_device()
            yield (self._cache_record.cached_model.get_cpu_state_dict(), self._cache_record.cached_model.model)
        finally:
            self._cache.unlock(self._cache_record)

    @property
    def model(self) -> AnyModel:
        """Return the model without locking it."""
        return self._cache_record.cached_model.model

    def repair_required_tensors_on_device(self) -> int:
        """Repair required tensors that should be resident on the cached model's execution device."""
        cached_model = self._cache_record.cached_model
        if not isinstance(cached_model, CachedModelWithPartialLoad):
            return 0
        return cached_model.repair_required_tensors_on_compute_device()


class LoadedModel(LoadedModelWithoutConfig):
    """Context manager object that mediates transfer from RAM<->VRAM."""

    def __init__(self, config: Optional[AnyModelConfig], cache_record: CacheRecord, cache: ModelCache):
        super().__init__(cache_record=cache_record, cache=cache)
        self.config = config


class ModelLoaderBase(ABC):
    """Abstract base class for loading models into RAM/VRAM."""

    @abstractmethod
    def __init__(
        self,
        app_config: InvokeAIAppConfig,
        logger: Logger,
        ram_cache: ModelCache,
    ):
        """Initialize the loader."""
        pass

    @abstractmethod
    def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
        """
        Return a model given its confguration.

        Given a model identified in the model configuration backend,
        return a ModelInfo object that can be used to retrieve the model.

        :param model_config: Model configuration, as returned by ModelConfigRecordStore
        :param submodel_type: an ModelType enum indicating the portion of
               the model to retrieve (e.g. ModelType.Vae)
        """
        pass

    @abstractmethod
    def get_size_fs(
        self, config: AnyModelConfig, model_path: Path, submodel_type: Optional[SubModelType] = None
    ) -> int:
        """Return size in bytes of the model, calculated before loading."""
        pass

    @property
    @abstractmethod
    def ram_cache(self) -> ModelCache:
        """Return the ram cache associated with this loader."""
        pass