diff --git a/docs/contributing/INVOCATIONS.md b/docs/contributing/INVOCATIONS.md index 1f12cfc8f5..212233f497 100644 --- a/docs/contributing/INVOCATIONS.md +++ b/docs/contributing/INVOCATIONS.md @@ -19,31 +19,56 @@ An invocation looks like this: ```py class UpscaleInvocation(BaseInvocation): """Upscales an image.""" - type: Literal['upscale'] = 'upscale' + + # fmt: off + type: Literal["upscale"] = "upscale" # Inputs - image: Union[ImageField,None] = Field(description="The input image") - strength: float = Field(default=0.75, gt=0, le=1, description="The strength") - level: Literal[2,4] = Field(default=2, description = "The upscale level") + image: Union[ImageField, None] = Field(description="The input image", default=None) + strength: float = Field(default=0.75, gt=0, le=1, description="The strength") + level: Literal[2, 4] = Field(default=2, description="The upscale level") + # fmt: on + + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["upscaling", "image"], + }, + } def invoke(self, context: InvocationContext) -> ImageOutput: - image = context.services.images.get(self.image.image_type, self.image.image_name) - results = context.services.generate.upscale_and_reconstruct( - image_list = [[image, 0]], - upscale = (self.level, self.strength), - strength = 0.0, # GFPGAN strength - save_original = False, - image_callback = None, + image = context.services.images.get_pil_image( + self.image.image_origin, self.image.image_name + ) + results = context.services.restoration.upscale_and_reconstruct( + image_list=[[image, 0]], + upscale=(self.level, self.strength), + strength=0.0, # GFPGAN strength + save_original=False, + image_callback=None, ) # Results are image and seed, unwrap for now # TODO: can this return multiple results? - image_type = ImageType.RESULT - image_name = context.services.images.create_name(context.graph_execution_state_id, self.id) - context.services.images.save(image_type, image_name, results[0][0]) - return ImageOutput( - image = ImageField(image_type = image_type, image_name = image_name) + image_dto = context.services.images.create( + image=results[0][0], + image_origin=ResourceOrigin.INTERNAL, + image_category=ImageCategory.GENERAL, + node_id=self.id, + session_id=context.graph_execution_state_id, + is_intermediate=self.is_intermediate, ) + + return ImageOutput( + image=ImageField( + image_name=image_dto.image_name, + image_origin=image_dto.image_origin, + ), + width=image_dto.width, + height=image_dto.height, + ) + ``` Each portion is important to implement correctly. @@ -95,25 +120,67 @@ Finally, note that for all linking, the `type` of the linked fields must match. If the `name` also matches, then the field can be **automatically linked** to a previous invocation by name and matching. +### Config + +```py + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["upscaling", "image"], + }, + } +``` + +This is an optional configuration for the invocation. It inherits from +pydantic's model `Config` class, and it used primarily to customize the +autogenerated OpenAPI schema. + +The UI relies on the OpenAPI schema in two ways: + +- An API client & Typescript types are generated from it. This happens at build + time. +- The node editor parses the schema into a template used by the UI to create the + node editor UI. This parsing happens at runtime. + +In this example, a `ui` key has been added to the `schema_extra` dict to provide +some tags for the UI, to facilitate filtering nodes. + +See the Schema Generation section below for more information. + ### Invoke Function ```py def invoke(self, context: InvocationContext) -> ImageOutput: - image = context.services.images.get(self.image.image_type, self.image.image_name) - results = context.services.generate.upscale_and_reconstruct( - image_list = [[image, 0]], - upscale = (self.level, self.strength), - strength = 0.0, # GFPGAN strength - save_original = False, - image_callback = None, + image = context.services.images.get_pil_image( + self.image.image_origin, self.image.image_name + ) + results = context.services.restoration.upscale_and_reconstruct( + image_list=[[image, 0]], + upscale=(self.level, self.strength), + strength=0.0, # GFPGAN strength + save_original=False, + image_callback=None, ) # Results are image and seed, unwrap for now - image_type = ImageType.RESULT - image_name = context.services.images.create_name(context.graph_execution_state_id, self.id) - context.services.images.save(image_type, image_name, results[0][0]) + # TODO: can this return multiple results? + image_dto = context.services.images.create( + image=results[0][0], + image_origin=ResourceOrigin.INTERNAL, + image_category=ImageCategory.GENERAL, + node_id=self.id, + session_id=context.graph_execution_state_id, + is_intermediate=self.is_intermediate, + ) + return ImageOutput( - image = ImageField(image_type = image_type, image_name = image_name) + image=ImageField( + image_name=image_dto.image_name, + image_origin=image_dto.image_origin, + ), + width=image_dto.width, + height=image_dto.height, ) ``` @@ -135,9 +202,16 @@ scenarios. If you need functionality, please provide it as a service in the ```py class ImageOutput(BaseInvocationOutput): """Base class for invocations that output an image""" - type: Literal['image'] = 'image' - image: ImageField = Field(default=None, description="The output image") + # fmt: off + type: Literal["image_output"] = "image_output" + image: ImageField = Field(default=None, description="The output image") + width: int = Field(description="The width of the image in pixels") + height: int = Field(description="The height of the image in pixels") + # fmt: on + + class Config: + schema_extra = {"required": ["type", "image", "width", "height"]} ``` Output classes look like an invocation class without the invoke method. Prefer @@ -168,35 +242,36 @@ Here's that `ImageOutput` class, without the needed schema customisation: class ImageOutput(BaseInvocationOutput): """Base class for invocations that output an image""" - type: Literal["image"] = "image" + # fmt: off + type: Literal["image_output"] = "image_output" image: ImageField = Field(default=None, description="The output image") + width: int = Field(description="The width of the image in pixels") + height: int = Field(description="The height of the image in pixels") + # fmt: on ``` -The generated OpenAPI schema, and all clients/types generated from it, will have -the `type` and `image` properties marked as optional, even though we know they -will always have a value by the time we can interact with them via the API. - -Here's the same class, but with the schema customisation added: +The OpenAPI schema that results from this `ImageOutput` will have the `type`, +`image`, `width` and `height` properties marked as optional, even though we know +they will always have a value. ```python class ImageOutput(BaseInvocationOutput): """Base class for invocations that output an image""" - type: Literal["image"] = "image" + # fmt: off + type: Literal["image_output"] = "image_output" image: ImageField = Field(default=None, description="The output image") + width: int = Field(description="The width of the image in pixels") + height: int = Field(description="The height of the image in pixels") + # fmt: on + # Add schema customization class Config: - schema_extra = { - 'required': [ - 'type', - 'image', - ] - } + schema_extra = {"required": ["type", "image", "width", "height"]} ``` -The resultant schema (and any API client or types generated from it) will now -have see `type` as string literal `"image"` and `image` as an `ImageField` -object. +With the customization in place, the schema will now show these properties as +required, obviating the need for extensive null checks in client code. See this `pydantic` issue for discussion on this solution: