mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
feat(nodes): make delete on startup configurable for obj serializer
- The default is to not delete on startup - feels safer. - The two services using this class _do_ delete on startup. - The class has "ephemeral" removed from its name. - Tests & app updated for this change.
This commit is contained in:
@ -5,7 +5,7 @@ from logging import Logger
|
||||
import torch
|
||||
|
||||
from invokeai.app.services.item_storage.item_storage_memory import ItemStorageMemory
|
||||
from invokeai.app.services.object_serializer.object_serializer_ephemeral_disk import ObjectSerializerEphemeralDisk
|
||||
from invokeai.app.services.object_serializer.object_serializer_disk import ObjectSerializerDisk
|
||||
from invokeai.app.services.object_serializer.object_serializer_forward_cache import ObjectSerializerForwardCache
|
||||
from invokeai.app.services.shared.sqlite.sqlite_util import init_db
|
||||
from invokeai.backend.model_manager.metadata import ModelMetadataStore
|
||||
@ -90,9 +90,11 @@ class ApiDependencies:
|
||||
image_records = SqliteImageRecordStorage(db=db)
|
||||
images = ImageService()
|
||||
invocation_cache = MemoryInvocationCache(max_cache_size=config.node_cache_size)
|
||||
tensors = ObjectSerializerForwardCache(ObjectSerializerEphemeralDisk[torch.Tensor](output_folder / "tensors"))
|
||||
tensors = ObjectSerializerForwardCache(
|
||||
ObjectSerializerDisk[torch.Tensor](output_folder / "tensors", delete_on_startup=True)
|
||||
)
|
||||
conditioning = ObjectSerializerForwardCache(
|
||||
ObjectSerializerEphemeralDisk[ConditioningFieldData](output_folder / "conditioning")
|
||||
ObjectSerializerDisk[ConditioningFieldData](output_folder / "conditioning", delete_on_startup=True)
|
||||
)
|
||||
model_manager = ModelManagerService(config, logger)
|
||||
model_record_service = ModelRecordServiceSQL(db=db)
|
||||
|
@ -22,26 +22,30 @@ class DeleteAllResult:
|
||||
freed_space_bytes: float
|
||||
|
||||
|
||||
class ObjectSerializerEphemeralDisk(ObjectSerializerBase[T]):
|
||||
"""Provides a disk-backed ephemeral storage for arbitrary python objects. The storage is cleared at startup.
|
||||
class ObjectSerializerDisk(ObjectSerializerBase[T]):
|
||||
"""Provides a disk-backed storage for arbitrary python objects.
|
||||
|
||||
:param output_folder: The folder where the objects will be stored
|
||||
:param delete_on_startup: If True, all objects in the output folder will be deleted on startup
|
||||
"""
|
||||
|
||||
def __init__(self, output_dir: Path):
|
||||
def __init__(self, output_dir: Path, delete_on_startup: bool = False):
|
||||
super().__init__()
|
||||
self._output_dir = output_dir
|
||||
self._output_dir.mkdir(parents=True, exist_ok=True)
|
||||
self._delete_on_startup = delete_on_startup
|
||||
self.__obj_class_name: Optional[str] = None
|
||||
|
||||
def start(self, invoker: "Invoker") -> None:
|
||||
self._invoker = invoker
|
||||
delete_all_result = self._delete_all()
|
||||
if delete_all_result.deleted_count > 0:
|
||||
freed_space_in_mb = round(delete_all_result.freed_space_bytes / 1024 / 1024, 2)
|
||||
self._invoker.services.logger.info(
|
||||
f"Deleted {delete_all_result.deleted_count} {self._obj_class_name} files (freed {freed_space_in_mb}MB)"
|
||||
)
|
||||
|
||||
if self._delete_on_startup:
|
||||
delete_all_result = self._delete_all()
|
||||
if delete_all_result.deleted_count > 0:
|
||||
freed_space_in_mb = round(delete_all_result.freed_space_bytes / 1024 / 1024, 2)
|
||||
self._invoker.services.logger.info(
|
||||
f"Deleted {delete_all_result.deleted_count} {self._obj_class_name} files (freed {freed_space_in_mb}MB)"
|
||||
)
|
||||
|
||||
def load(self, name: str) -> T:
|
||||
file_path = self._get_path(name)
|
Reference in New Issue
Block a user