From 33a0af46375c38d28c576d03ef18835729fc8f52 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Sat, 27 May 2023 09:10:02 +1000 Subject: [PATCH] feat(nodes): add nameservice Currenly only used to make names for images, but when latents, conditioning, etc are managed in DB, will do the same for them. Intended to eventually support custom naming schemes. --- invokeai/app/api/dependencies.py | 4 ++- invokeai/app/cli_app.py | 3 ++ invokeai/app/services/image_record_storage.py | 1 - invokeai/app/services/images.py | 30 +++++-------------- invokeai/app/services/resource_name.py | 30 +++++++++++++++++++ 5 files changed, 44 insertions(+), 24 deletions(-) create mode 100644 invokeai/app/services/resource_name.py diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py index dfef5d2176..0dfb5505b6 100644 --- a/invokeai/app/api/dependencies.py +++ b/invokeai/app/api/dependencies.py @@ -5,6 +5,7 @@ import os from invokeai.app.services.image_record_storage import SqliteImageRecordStorage from invokeai.app.services.images import ImageService from invokeai.app.services.metadata import CoreMetadataService +from invokeai.app.services.resource_name import SimpleNameService from invokeai.app.services.urls import LocalUrlService from invokeai.backend.util.logging import InvokeAILogger @@ -67,7 +68,7 @@ class ApiDependencies: metadata = CoreMetadataService() image_record_storage = SqliteImageRecordStorage(db_location) image_file_storage = DiskImageFileStorage(f"{output_folder}/images") - + names = SimpleNameService() latents = ForwardCacheLatentsStorage( DiskLatentsStorage(f"{output_folder}/latents") ) @@ -78,6 +79,7 @@ class ApiDependencies: metadata=metadata, url=urls, logger=logger, + names=names, graph_execution_manager=graph_execution_manager, ) diff --git a/invokeai/app/cli_app.py b/invokeai/app/cli_app.py index de543d2d85..eb55ba45d2 100644 --- a/invokeai/app/cli_app.py +++ b/invokeai/app/cli_app.py @@ -16,6 +16,7 @@ from pydantic.fields import Field from invokeai.app.services.image_record_storage import SqliteImageRecordStorage from invokeai.app.services.images import ImageService from invokeai.app.services.metadata import CoreMetadataService +from invokeai.app.services.resource_name import SimpleNameService from invokeai.app.services.urls import LocalUrlService @@ -229,6 +230,7 @@ def invoke_cli(): metadata = CoreMetadataService() image_record_storage = SqliteImageRecordStorage(db_location) image_file_storage = DiskImageFileStorage(f"{output_folder}/images") + names = SimpleNameService() images = ImageService( image_record_storage=image_record_storage, @@ -236,6 +238,7 @@ def invoke_cli(): metadata=metadata, url=urls, logger=logger, + names=names, graph_execution_manager=graph_execution_manager, ) diff --git a/invokeai/app/services/image_record_storage.py b/invokeai/app/services/image_record_storage.py index 188a411a6b..9a73b68e21 100644 --- a/invokeai/app/services/image_record_storage.py +++ b/invokeai/app/services/image_record_storage.py @@ -103,7 +103,6 @@ class SqliteImageRecordStorage(ImageRecordStorageBase): def __init__(self, filename: str) -> None: super().__init__() - self._filename = filename self._conn = sqlite3.connect(filename, check_same_thread=False) # Enable row factory to get rows as dictionaries (must be done before making the cursor!) diff --git a/invokeai/app/services/images.py b/invokeai/app/services/images.py index d0f7236fe2..6da7510702 100644 --- a/invokeai/app/services/images.py +++ b/invokeai/app/services/images.py @@ -1,5 +1,6 @@ from abc import ABC, abstractmethod from logging import Logger +from os import name from typing import Optional, TYPE_CHECKING, Union import uuid from PIL.Image import Image as PILImageType @@ -31,6 +32,7 @@ from invokeai.app.services.image_file_storage import ( ) from invokeai.app.services.item_storage import ItemStorageABC, PaginatedResults from invokeai.app.services.metadata import MetadataServiceBase +from invokeai.app.services.resource_name import NameServiceBase from invokeai.app.services.urls import UrlServiceBase if TYPE_CHECKING: @@ -120,6 +122,7 @@ class ImageServiceDependencies: metadata: MetadataServiceBase urls: UrlServiceBase logger: Logger + names: NameServiceBase graph_execution_manager: ItemStorageABC["GraphExecutionState"] def __init__( @@ -129,6 +132,7 @@ class ImageServiceDependencies: metadata: MetadataServiceBase, url: UrlServiceBase, logger: Logger, + names: NameServiceBase, graph_execution_manager: ItemStorageABC["GraphExecutionState"], ): self.records = image_record_storage @@ -136,6 +140,7 @@ class ImageServiceDependencies: self.metadata = metadata self.urls = url self.logger = logger + self.names = names self.graph_execution_manager = graph_execution_manager @@ -149,6 +154,7 @@ class ImageService(ImageServiceABC): metadata: MetadataServiceBase, url: UrlServiceBase, logger: Logger, + names: NameServiceBase, graph_execution_manager: ItemStorageABC["GraphExecutionState"], ): self._services = ImageServiceDependencies( @@ -157,6 +163,7 @@ class ImageService(ImageServiceABC): metadata=metadata, url=url, logger=logger, + names=names, graph_execution_manager=graph_execution_manager, ) @@ -175,12 +182,7 @@ class ImageService(ImageServiceABC): if image_category not in ImageCategory: raise InvalidImageCategoryException - image_name = self._create_image_name( - image_type=image_type, - image_category=image_category, - node_id=node_id, - session_id=session_id, - ) + image_name = self._services.names.create_image_name() metadata = self._get_metadata(session_id, node_id) @@ -260,7 +262,6 @@ class ImageService(ImageServiceABC): except Exception as e: self._services.logger.error("Problem updating image record") raise e - def get_pil_image(self, image_type: ImageType, image_name: str) -> PILImageType: try: @@ -378,21 +379,6 @@ class ImageService(ImageServiceABC): self._services.logger.error("Problem deleting image record and file") raise e - def _create_image_name( - self, - image_type: ImageType, - image_category: ImageCategory, - node_id: Optional[str] = None, - session_id: Optional[str] = None, - ) -> str: - """Create a unique image name.""" - uuid_str = str(uuid.uuid4()) - - if node_id is not None and session_id is not None: - return f"{image_type.value}_{image_category.value}_{session_id}_{node_id}_{uuid_str}.png" - - return f"{image_type.value}_{image_category.value}_{uuid_str}.png" - def _get_metadata( self, session_id: Optional[str] = None, node_id: Optional[str] = None ) -> Union[ImageMetadata, None]: diff --git a/invokeai/app/services/resource_name.py b/invokeai/app/services/resource_name.py new file mode 100644 index 0000000000..dd5a76cfc0 --- /dev/null +++ b/invokeai/app/services/resource_name.py @@ -0,0 +1,30 @@ +from abc import ABC, abstractmethod +from enum import Enum, EnumMeta +import uuid + + +class ResourceType(str, Enum, metaclass=EnumMeta): + """Enum for resource types.""" + + IMAGE = "image" + LATENT = "latent" + + +class NameServiceBase(ABC): + """Low-level service responsible for naming resources (images, latents, etc).""" + + # TODO: Add customizable naming schemes + @abstractmethod + def create_image_name(self) -> str: + """Creates a name for an image.""" + pass + + +class SimpleNameService(NameServiceBase): + """Creates image names from UUIDs.""" + + # TODO: Add customizable naming schemes + def create_image_name(self) -> str: + uuid_str = str(uuid.uuid4()) + filename = f"{uuid_str}.png" + return filename