mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
07e3a0ec15
- 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.
57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
|
|
|
|
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 .image import ImageOutput
|
|
|
|
|
|
class UpscaleInvocation(BaseInvocation):
|
|
"""Upscales an image."""
|
|
#fmt: off
|
|
type: Literal["upscale"] = "upscale"
|
|
|
|
# Inputs
|
|
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.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)
|
|
)
|