From 3e5a91e3bff00a1189409dc8ad8085aedf4b00f4 Mon Sep 17 00:00:00 2001 From: dunkeroni Date: Sun, 18 Feb 2024 16:56:46 -0500 Subject: [PATCH] chore(invocations): use IMAGE_MODES constant literal --- invokeai/app/invocations/constants.py | 3 +++ invokeai/app/invocations/image.py | 4 +--- invokeai/app/services/shared/invocation_context.py | 11 ++++++----- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/invokeai/app/invocations/constants.py b/invokeai/app/invocations/constants.py index 795e7a3b60..fca5a2ec7f 100644 --- a/invokeai/app/invocations/constants.py +++ b/invokeai/app/invocations/constants.py @@ -12,3 +12,6 @@ The ratio of image:latent dimensions is LATENT_SCALE_FACTOR:1, or 8:1. SCHEDULER_NAME_VALUES = Literal[tuple(SCHEDULER_MAP.keys())] """A literal type representing the valid scheduler names.""" + +IMAGE_MODES = Literal["L", "RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV", "I", "F"] +"""A literal type for PIL image modes supported by Invoke""" \ No newline at end of file diff --git a/invokeai/app/invocations/image.py b/invokeai/app/invocations/image.py index f5ad5515a6..1f3b5b7368 100644 --- a/invokeai/app/invocations/image.py +++ b/invokeai/app/invocations/image.py @@ -16,6 +16,7 @@ from invokeai.app.invocations.fields import ( WithMetadata, ) from invokeai.app.invocations.primitives import ImageOutput +from invokeai.app.invocations.constants import IMAGE_MODES from invokeai.app.services.image_records.image_records_common import ImageCategory from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.image_util.invisible_watermark import InvisibleWatermark @@ -263,9 +264,6 @@ class ImageChannelInvocation(BaseInvocation, WithMetadata, WithBoard): return ImageOutput.build(image_dto) -IMAGE_MODES = Literal["L", "RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV", "I", "F"] - - @invocation( "img_conv", title="Convert Image Mode", diff --git a/invokeai/app/services/shared/invocation_context.py b/invokeai/app/services/shared/invocation_context.py index d217a865af..2383785ad4 100644 --- a/invokeai/app/services/shared/invocation_context.py +++ b/invokeai/app/services/shared/invocation_context.py @@ -6,6 +6,7 @@ from PIL.Image import Image from torch import Tensor from invokeai.app.invocations.fields import MetadataField, WithBoard, WithMetadata +from invokeai.app.invocations.constants import IMAGE_MODES from invokeai.app.services.boards.boards_common import BoardDTO from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin @@ -194,19 +195,19 @@ class ImagesInterface(InvocationContextInterface): node_id=self._context_data.invocation.id, ) - def get_pil(self, image_name: str, format: str | None = None) -> Image: + def get_pil(self, image_name: str, mode: IMAGE_MODES | None = None) -> Image: """ Gets an image as a PIL Image object. :param image_name: The name of the image to get. - :param format: The color format to convert the image to. If None, the original format is used. + :param mode: The color mode to convert the image to. If None, the original mode is used. """ image = self._services.images.get_pil_image(image_name) - if format and format != image.mode: + if mode and mode != image.mode: try: - image = image.convert(format) + image = image.convert(mode) except ValueError: - self._services.logger.warning(f"Could not convert image from {image.mode} to {format}. Using original format.") + self._services.logger.warning(f"Could not convert image from {image.mode} to {mode}. Using original mode instead.") return image def get_metadata(self, image_name: str) -> Optional[MetadataField]: