diff --git a/invokeai/app/invocations/baseinvocation.py b/invokeai/app/invocations/baseinvocation.py index 72fe39ed0b..3590129b96 100644 --- a/invokeai/app/invocations/baseinvocation.py +++ b/invokeai/app/invocations/baseinvocation.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from inspect import signature -from typing import get_args, get_type_hints +from typing import get_args, get_type_hints, Dict, List, Literal, TypedDict from pydantic import BaseModel, Field @@ -76,3 +76,56 @@ class BaseInvocation(ABC, BaseModel): #fmt: off id: str = Field(description="The id of this node. Must be unique among all nodes.") #fmt: on + + +# TODO: figure out a better way to provide these hints +# TODO: when we can upgrade to python 3.11, we can use the`NotRequired` type instead of `total=False` +class UIConfig(TypedDict, total=False): + type_hints: Dict[ + str, + Literal[ + "integer", + "float", + "boolean", + "string", + "enum", + "image", + "latents", + "model", + ], + ] + tags: List[str] + + +class CustomisedSchemaExtra(TypedDict): + ui: UIConfig + + +class InvocationConfig(BaseModel.Config): + """Customizes pydantic's BaseModel.Config class for use by Invocations. + + Provide `schema_extra` a `ui` dict to add hints for generated UIs. + + `tags` + - A list of strings, used to categorise invocations. + + `type_hints` + - A dict of field types which override the types in the invocation definition. + - Each key should be the name of one of the invocation's fields. + - Each value should be one of the valid types: + - `integer`, `float`, `boolean`, `string`, `enum`, `image`, `latents`, `model` + + ```python + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["stable-diffusion", "image"], + "type_hints": { + "initial_image": "image", + }, + }, + } + ``` + """ + + schema_extra: CustomisedSchemaExtra diff --git a/invokeai/app/invocations/cv.py b/invokeai/app/invocations/cv.py index adcb2405c4..9afbbbbcc9 100644 --- a/invokeai/app/invocations/cv.py +++ b/invokeai/app/invocations/cv.py @@ -6,10 +6,9 @@ import cv2 as cv import numpy from PIL import Image, ImageOps from pydantic import BaseModel, Field -from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.image import ImageField, ImageType -from .baseinvocation import BaseInvocation, InvocationContext +from .baseinvocation import BaseInvocation, InvocationContext, InvocationConfig from .image import ImageOutput diff --git a/invokeai/app/invocations/generate.py b/invokeai/app/invocations/generate.py index a07b1ac379..d0eeeae698 100644 --- a/invokeai/app/invocations/generate.py +++ b/invokeai/app/invocations/generate.py @@ -7,13 +7,10 @@ import numpy as np from torch import Tensor from pydantic import BaseModel, Field -from invokeai.app.invocations.models.config import ( - InvocationConfig, -) from invokeai.app.models.image import ImageField, ImageType from invokeai.app.invocations.util.get_model import choose_model -from .baseinvocation import BaseInvocation, InvocationContext +from .baseinvocation import BaseInvocation, InvocationContext, InvocationConfig from .image import ImageOutput from ...backend.generator import Txt2Img, Img2Img, Inpaint, InvokeAIGenerator from ...backend.stable_diffusion import PipelineIntermediateState diff --git a/invokeai/app/invocations/image.py b/invokeai/app/invocations/image.py index 0f783b2541..cc5f6b53c7 100644 --- a/invokeai/app/invocations/image.py +++ b/invokeai/app/invocations/image.py @@ -7,11 +7,9 @@ import numpy from PIL import Image, ImageFilter, ImageOps from pydantic import BaseModel, Field -from invokeai.app.invocations.models.config import InvocationConfig - from ..models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices -from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext +from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig class PILInvocationConfig(BaseModel): diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py index d59d681f40..2da6e451a9 100644 --- a/invokeai/app/invocations/latent.py +++ b/invokeai/app/invocations/latent.py @@ -4,7 +4,6 @@ from typing import Literal, Optional from pydantic import BaseModel, Field import torch -from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.exceptions import CanceledException from invokeai.app.invocations.util.get_model import choose_model from invokeai.app.util.step_callback import diffusers_step_callback_adapter @@ -15,7 +14,7 @@ from ...backend.stable_diffusion.diffusion.shared_invokeai_diffusion import Post from ...backend.image_util.seamless import configure_model_padding from ...backend.prompting.conditioning import get_uc_and_c_and_ec from ...backend.stable_diffusion.diffusers_pipeline import ConditioningData, StableDiffusionGeneratorPipeline -from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext +from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig import numpy as np from ..services.image_storage import ImageType from .baseinvocation import BaseInvocation, InvocationContext diff --git a/invokeai/app/invocations/math.py b/invokeai/app/invocations/math.py index e9d90b1d61..afb0e75377 100644 --- a/invokeai/app/invocations/math.py +++ b/invokeai/app/invocations/math.py @@ -4,9 +4,7 @@ from typing import Literal from pydantic import BaseModel, Field -from invokeai.app.invocations.models.config import InvocationConfig - -from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext +from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig class MathInvocationConfig(BaseModel): diff --git a/invokeai/app/invocations/models/config.py b/invokeai/app/invocations/models/config.py deleted file mode 100644 index f53bbdda00..0000000000 --- a/invokeai/app/invocations/models/config.py +++ /dev/null @@ -1,54 +0,0 @@ -from typing import Dict, List, Literal, TypedDict -from pydantic import BaseModel - - -# TODO: when we can upgrade to python 3.11, we can use the`NotRequired` type instead of `total=False` -class UIConfig(TypedDict, total=False): - type_hints: Dict[ - str, - Literal[ - "integer", - "float", - "boolean", - "string", - "enum", - "image", - "latents", - "model", - ], - ] - tags: List[str] - - -class CustomisedSchemaExtra(TypedDict): - ui: UIConfig - - -class InvocationConfig(BaseModel.Config): - """Customizes pydantic's BaseModel.Config class for use by Invocations. - - Provide `schema_extra` a `ui` dict to add hints for generated UIs. - - `tags` - - A list of strings, used to categorise invocations. - - `type_hints` - - A dict of field types which override the types in the invocation definition. - - Each key should be the name of one of the invocation's fields. - - Each value should be one of the valid types: - - `integer`, `float`, `boolean`, `string`, `enum`, `image`, `latents`, `model` - - ```python - class Config(InvocationConfig): - schema_extra = { - "ui": { - "tags": ["stable-diffusion", "image"], - "type_hints": { - "initial_image": "image", - }, - }, - } - ``` - """ - - schema_extra: CustomisedSchemaExtra diff --git a/invokeai/app/invocations/reconstruct.py b/invokeai/app/invocations/reconstruct.py index db0a28c075..f6df5a2254 100644 --- a/invokeai/app/invocations/reconstruct.py +++ b/invokeai/app/invocations/reconstruct.py @@ -2,11 +2,10 @@ from datetime import datetime, timezone from typing import Literal, Union from pydantic import Field -from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices -from .baseinvocation import BaseInvocation, InvocationContext +from .baseinvocation import BaseInvocation, InvocationContext, InvocationConfig from .image import ImageOutput class RestoreFaceInvocation(BaseInvocation): diff --git a/invokeai/app/invocations/upscale.py b/invokeai/app/invocations/upscale.py index c9aa86cc6d..021f3569e8 100644 --- a/invokeai/app/invocations/upscale.py +++ b/invokeai/app/invocations/upscale.py @@ -4,11 +4,10 @@ from datetime import datetime, timezone from typing import Literal, Union from pydantic import Field -from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices -from .baseinvocation import BaseInvocation, InvocationContext +from .baseinvocation import BaseInvocation, InvocationContext, InvocationConfig from .image import ImageOutput