Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add some docs to OriginalWeightsStorage and fix type hints.

Browse Source
This commit is contained in:
Ryan Dick 2024-07-31 15:08:24 -04:00
parent 86f705bf48
commit 0bb7ed44f6
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
invokeai/backend/util/original_weights_storage.py
View File

@ -8,9 +8,13 @@ from invokeai.backend.util.devices import TorchDevice
class OriginalWeightsStorage: class OriginalWeightsStorage:
"""A class for tracking the original weights of a model for patch/unpatch operations."""
def __init__(self, cached_weights: Optional[Dict[str, torch.Tensor]] = None): def __init__(self, cached_weights: Optional[Dict[str, torch.Tensor]] = None):
self._weights = {} # The original weights of the model.
self._changed_weights = set() self._weights: dict[str, torch.Tensor] = {}
# The keys of the weights that have been changed (via `save()`) during the lifetime of this instance.
self._changed_weights: set[str] = set()
if cached_weights: if cached_weights:
self._weights.update(cached_weights) self._weights.update(cached_weights)