diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index 9b88f82fa3..0ab1ac241a 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -72,6 +72,7 @@ example_model_input = { "description": "Model description", "vae": None, "variant": "normal", + "image": "blob" } ############################################################################## diff --git a/invokeai/app/services/model_images/model_images_base.py b/invokeai/app/services/model_images/model_images_base.py new file mode 100644 index 0000000000..624cdd8216 --- /dev/null +++ b/invokeai/app/services/model_images/model_images_base.py @@ -0,0 +1,32 @@ +from abc import ABC, abstractmethod +from pathlib import Path + +from PIL.Image import Image as PILImageType + + +class ModelImagesBase(ABC): + """Low-level service responsible for storing and retrieving image files.""" + + @abstractmethod + def get(self, image_name: str) -> PILImageType: + """Retrieves an image as PIL Image.""" + pass + + @abstractmethod + def get_path(self, image_name: str) -> Path: + """Gets the internal path to an image.""" + pass + + @abstractmethod + def save( + self, + image: PILImageType, + image_name: str, + ) -> None: + """Saves an image. Returns a tuple of the image name and created timestamp.""" + pass + + @abstractmethod + def delete(self, image_name: str) -> None: + """Deletes an image.""" + pass diff --git a/invokeai/app/services/model_images/model_images_common.py b/invokeai/app/services/model_images/model_images_common.py new file mode 100644 index 0000000000..e63b780261 --- /dev/null +++ b/invokeai/app/services/model_images/model_images_common.py @@ -0,0 +1,20 @@ +# TODO: Should these excpetions subclass existing python exceptions? +class ModelImageFileNotFoundException(Exception): + """Raised when an image file is not found in storage.""" + + def __init__(self, message="Image file not found"): + super().__init__(message) + + +class ModelImageFileSaveException(Exception): + """Raised when an image cannot be saved.""" + + def __init__(self, message="Image file not saved"): + super().__init__(message) + + +class ModelImageFileDeleteException(Exception): + """Raised when an image cannot be deleted.""" + + def __init__(self, message="Image file not deleted"): + super().__init__(message) diff --git a/invokeai/app/services/model_images/model_images_disk.py b/invokeai/app/services/model_images/model_images_disk.py new file mode 100644 index 0000000000..340a3d6b6f --- /dev/null +++ b/invokeai/app/services/model_images/model_images_disk.py @@ -0,0 +1,82 @@ +from pathlib import Path +from typing import Union + +from PIL import Image, PngImagePlugin +from PIL.Image import Image as PILImageType +from send2trash import send2trash + +from invokeai.app.services.invoker import Invoker + +from .model_images_base import ModelImagesBase +from .model_images_common import ModelImageFileDeleteException, ModelImageFileNotFoundException, ModelImageFileSaveException + + +class DiskImageFileStorage(ModelImagesBase): + """Stores images on disk""" + + __output_folder: Path + __invoker: Invoker + + def __init__(self, output_folder: Union[str, Path]): + + self.__output_folder: Path = output_folder if isinstance(output_folder, Path) else Path(output_folder) + # Validate required output folders at launch + self.__validate_storage_folders() + + def start(self, invoker: Invoker) -> None: + self.__invoker = invoker + + def get(self, image_name: str) -> PILImageType: + try: + image_path = self.get_path(image_name) + + image = Image.open(image_path) + return image + except FileNotFoundError as e: + raise ModelImageFileNotFoundException from e + + def save( + self, + image: PILImageType, + image_name: str, + ) -> None: + try: + self.__validate_storage_folders() + image_path = self.get_path(image_name) + + pnginfo = PngImagePlugin.PngInfo() + info_dict = {} + + # When saving the image, the image object's info field is not populated. We need to set it + image.info = info_dict + image.save( + image_path, + "PNG", + pnginfo=pnginfo, + compress_level=self.__invoker.services.configuration.png_compress_level, + ) + + except Exception as e: + raise ModelImageFileSaveException from e + + def delete(self, image_name: str) -> None: + try: + image_path = self.get_path(image_name) + + if image_path.exists(): + send2trash(image_path) + + except Exception as e: + raise ModelImageFileDeleteException from e + + # TODO: make this a bit more flexible for e.g. cloud storage + def get_path(self, image_name: str) -> Path: + path = self.__output_folder / image_name + + return path + + def __validate_storage_folders(self) -> None: + """Checks if the required output folders exist and create them if they don't""" + folders: list[Path] = [self.__output_folder] + for folder in folders: + folder.mkdir(parents=True, exist_ok=True) diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelEdit.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelEdit.tsx index 54b6204e8e..6be73ffe28 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelEdit.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelEdit.tsx @@ -23,6 +23,7 @@ import type { UpdateModelArg } from 'services/api/endpoints/models'; import { useGetModelConfigQuery, useUpdateModelMutation } from 'services/api/endpoints/models'; import BaseModelSelect from './Fields/BaseModelSelect'; +import ModelImageUpload from './Fields/ModelImageUpload'; import ModelVariantSelect from './Fields/ModelVariantSelect'; import PredictionTypeSelect from './Fields/PredictionTypeSelect'; @@ -58,6 +59,7 @@ export const ModelEdit = () => { key: data.key, body: values, }; + console.log(responseBody, 'responseBody') updateModel(responseBody) .unwrap() @@ -129,7 +131,8 @@ export const ModelEdit = () => { - + + {t('modelManager.description')}