feat(nodes): add invocation schema customisation, add model selection

- add invocation schema customisation

done via fastapi's `Config` class and `schema_extra`. when using `Config`, inherit from `InvocationConfig` to get type hints.

where it makes sense - like for all math invocations - define a `MathInvocationConfig` class and have all invocations inherit from it.

this customisation can provide any arbitrary additional data to the UI. currently it provides tags and field type hints.

this is necessary for `model` type fields, which are actually string fields. without something like this, we can't reliably differentiate  `model` fields from normal `string` fields.

can also be used for future field types.

all invocations now have tags, and all `model` fields have ui type hints.

- fix model handling for invocations

added a helper to fall back to the default model if an invalid model name is chosen. model names in graphs now work.

- fix latents progress callback

noticed this wasn't correct while working on everything else.
This commit is contained in:
psychedelicious
2023-04-10 19:07:48 +10:00
parent 427db7c7e2
commit 07e3a0ec15
9 changed files with 228 additions and 64 deletions

View File

@ -1,17 +1,26 @@
# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
from datetime import datetime, timezone
from typing import Literal, Optional
from typing import Literal
import numpy
from PIL import Image, ImageFilter, ImageOps
from pydantic import BaseModel, Field
from ..services.image_storage import ImageType
from ..services.invocation_services import InvocationServices
from invokeai.app.invocations.models.config import InvocationConfig
from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext
class MathInvocationConfig(BaseModel):
"""Helper class to provide all math invocations with additional config"""
# Schema customisation
class Config(InvocationConfig):
schema_extra = {
"ui": {
"tags": ["math"],
}
}
class IntOutput(BaseInvocationOutput):
"""An integer output"""
#fmt: off
@ -20,7 +29,7 @@ class IntOutput(BaseInvocationOutput):
#fmt: on
class AddInvocation(BaseInvocation):
class AddInvocation(BaseInvocation, MathInvocationConfig):
"""Adds two numbers"""
#fmt: off
type: Literal["add"] = "add"
@ -32,7 +41,7 @@ class AddInvocation(BaseInvocation):
return IntOutput(a=self.a + self.b)
class SubtractInvocation(BaseInvocation):
class SubtractInvocation(BaseInvocation, MathInvocationConfig):
"""Subtracts two numbers"""
#fmt: off
type: Literal["sub"] = "sub"
@ -44,7 +53,7 @@ class SubtractInvocation(BaseInvocation):
return IntOutput(a=self.a - self.b)
class MultiplyInvocation(BaseInvocation):
class MultiplyInvocation(BaseInvocation, MathInvocationConfig):
"""Multiplies two numbers"""
#fmt: off
type: Literal["mul"] = "mul"
@ -56,7 +65,7 @@ class MultiplyInvocation(BaseInvocation):
return IntOutput(a=self.a * self.b)
class DivideInvocation(BaseInvocation):
class DivideInvocation(BaseInvocation, MathInvocationConfig):
"""Divides two numbers"""
#fmt: off
type: Literal["div"] = "div"