DiffusionAi/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2025-01-08 11:57:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

Rip out ModelLockerBase.

Browse Source
This commit is contained in:
Ryan Dick 2024-12-04 21:47:11 +00:00
parent 712674b6dd
commit e48dee4c4a
6 changed files with 12 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
docs/contributing/MODEL_MANAGER.md
View File

@ -1364,7 +1364,7 @@ the in-memory loaded model:
|----------------|-----------------|------------------|
| `config` | AnyModelConfig | A copy of the model's configuration record for retrieving base type, etc. |
| `model` | AnyModel | The instantiated model (details below) |
| `locker` | ModelLockerBase | A context manager that mediates the movement of the model into VRAM |
| `locker` | ModelLocker | A context manager that mediates the movement of the model into VRAM |
### get_model_by_key(key, [submodel]) -> LoadedModel

5
invokeai/backend/model_manager/load/load_base.py
View File

@ -18,7 +18,8 @@ from invokeai.backend.model_manager.config import (
AnyModelConfig,
SubModelType,
)
from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase, ModelLockerBase
from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
from invokeai.backend.model_manager.load.model_cache.model_locker import ModelLocker
@dataclass
@ -60,7 +61,7 @@ class LoadedModelWithoutConfig:
not have a state_dict, in which case this value will be None.
"""
_locker: ModelLockerBase
_locker: ModelLocker
def __enter__(self) -> AnyModel:
"""Context entry."""

5
invokeai/backend/model_manager/load/load_default.py
View File

@ -14,7 +14,8 @@ from invokeai.backend.model_manager import (
)
from invokeai.backend.model_manager.config import DiffusersConfigBase
from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoaderBase
from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase, ModelLockerBase
from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
from invokeai.backend.model_manager.load.model_cache.model_locker import ModelLocker
from invokeai.backend.model_manager.load.model_util import calc_model_size_by_fs
from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
from invokeai.backend.util.devices import TorchDevice
@ -66,7 +67,7 @@ class ModelLoader(ModelLoaderBase):
model_base = self._app_config.models_path
return (model_base / config.path).resolve()
def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> ModelLockerBase:
def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> ModelLocker:
stats_name = ":".join([config.base, config.type, config.name, (submodel_type or "")])
try:
return self._ram_cache.get(config.key, submodel_type, stats_name=stats_name)

29
invokeai/backend/model_manager/load/model_cache/model_cache_base.py
View File

@ -15,32 +15,7 @@ from typing import Dict, Generic, Optional, TypeVar
import torch
from invokeai.backend.model_manager.config import AnyModel, SubModelType
class ModelLockerBase(ABC):
"""Base class for the model locker used by the loader."""
@abstractmethod
def lock(self) -> AnyModel:
"""Lock the contained model and move it into VRAM."""
pass
@abstractmethod
def unlock(self) -> None:
"""Unlock the contained model, and remove it from VRAM."""
pass
@abstractmethod
def get_state_dict(self) -> Optional[Dict[str, torch.Tensor]]:
"""Return the state dict (if any) for the cached model."""
pass
@property
@abstractmethod
def model(self) -> AnyModel:
"""Return the model."""
pass
from invokeai.backend.model_manager.load.model_cache.model_locker import ModelLocker
T = TypeVar("T")
@ -197,7 +172,7 @@ class ModelCacheBase(ABC, Generic[T]):
key: str,
submodel_type: Optional[SubModelType] = None,
stats_name: Optional[str] = None,
) -> ModelLockerBase:
) -> ModelLocker:
"""
Retrieve model using key and optional submodel_type.

3
invokeai/backend/model_manager/load/model_cache/model_cache_default.py
View File

@ -17,7 +17,6 @@ from invokeai.backend.model_manager.load.model_cache.model_cache_base import (
CacheRecord,
CacheStats,
ModelCacheBase,
ModelLockerBase,
)
from invokeai.backend.model_manager.load.model_cache.model_locker import ModelLocker
from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data
@ -184,7 +183,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
key: str,
submodel_type: Optional[SubModelType] = None,
stats_name: Optional[str] = None,
) -> ModelLockerBase:
) -> ModelLocker:
"""
Retrieve model using key and optional submodel_type.

8
invokeai/backend/model_manager/load/model_cache/model_locker.py
View File

@ -7,14 +7,10 @@ from typing import Dict, Optional
import torch
from invokeai.backend.model_manager import AnyModel
from invokeai.backend.model_manager.load.model_cache.model_cache_base import (
CacheRecord,
ModelCacheBase,
ModelLockerBase,
)
from invokeai.backend.model_manager.load.model_cache.model_cache_base import CacheRecord, ModelCacheBase
class ModelLocker(ModelLockerBase):
class ModelLocker:
"""Internal class that mediates movement in and out of GPU."""
def __init__(self, cache: ModelCacheBase[AnyModel], cache_entry: CacheRecord[AnyModel]):