mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2025-04-04 22:43:40 +08:00
2871676f79
* allow model patcher to optimize away the unpatching step when feasible * remove lazy_offloading functionality * allow model patcher to optimize away the unpatching step when feasible * remove lazy_offloading functionality * do not save original weights if there is a CPU copy of state dict * Update invokeai/backend/model_manager/load/load_base.py Co-authored-by: Ryan Dick <ryanjdick3@gmail.com> * documentation fixes added during penultimate review --------- Co-authored-by: Lincoln Stein <lstein@gmail.com> Co-authored-by: Kent Keirsey <31807370+hipsterusername@users.noreply.github.com> Co-authored-by: Ryan Dick <ryanjdick3@gmail.com>
69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
"""
|
|
Base class and implementation of a class that moves models in and out of VRAM.
|
|
"""
|
|
|
|
from typing import Dict, Optional
|
|
|
|
import torch
|
|
|
|
from invokeai.backend.model_manager import AnyModel
|
|
|
|
from .model_cache_base import CacheRecord, ModelCacheBase, ModelLockerBase
|
|
|
|
|
|
class ModelLocker(ModelLockerBase):
|
|
"""Internal class that mediates movement in and out of GPU."""
|
|
|
|
def __init__(self, cache: ModelCacheBase[AnyModel], cache_entry: CacheRecord[AnyModel]):
|
|
"""
|
|
Initialize the model locker.
|
|
|
|
:param cache: The ModelCache object
|
|
:param cache_entry: The entry in the model cache
|
|
"""
|
|
self._cache = cache
|
|
self._cache_entry = cache_entry
|
|
|
|
@property
|
|
def model(self) -> AnyModel:
|
|
"""Return the model without moving it around."""
|
|
return self._cache_entry.model
|
|
|
|
def get_state_dict(self) -> Optional[Dict[str, torch.Tensor]]:
|
|
"""Return the state dict (if any) for the cached model."""
|
|
return self._cache_entry.state_dict
|
|
|
|
def lock(self) -> AnyModel:
|
|
"""Move the model into the execution device (GPU) and lock it."""
|
|
if not hasattr(self.model, "to"):
|
|
return self.model
|
|
|
|
# NOTE that the model has to have the to() method in order for this code to move it into GPU!
|
|
self._cache_entry.lock()
|
|
try:
|
|
if self._cache.lazy_offloading:
|
|
self._cache.offload_unlocked_models(self._cache_entry.size)
|
|
self._cache.move_model_to_device(self._cache_entry, self._cache.execution_device)
|
|
self._cache_entry.loaded = True
|
|
self._cache.logger.debug(f"Locking {self._cache_entry.key} in {self._cache.execution_device}")
|
|
self._cache.print_cuda_stats()
|
|
except torch.cuda.OutOfMemoryError:
|
|
self._cache.logger.warning("Insufficient GPU memory to load model. Aborting")
|
|
self._cache_entry.unlock()
|
|
raise
|
|
except Exception:
|
|
self._cache_entry.unlock()
|
|
raise
|
|
|
|
return self.model
|
|
|
|
def unlock(self) -> None:
|
|
"""Call upon exit from context."""
|
|
if not hasattr(self.model, "to"):
|
|
return
|
|
|
|
self._cache_entry.unlock()
|
|
if not self._cache.lazy_offloading:
|
|
self._cache.offload_unlocked_models(0)
|
|
self._cache.print_cuda_stats()
|