mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
a1773197e9
- remove `image_origin` from most places where we interact with images - consolidate image file storage into a single `images/` dir Images have an `image_origin` attribute but it is not actually used when retrieving images, nor will it ever be. It is still used when creating images and helps to differentiate between internally generated images and uploads. It was included in eg API routes and image service methods as a holdover from the previous app implementation where images were not managed in a database. Now that we have images in a db, we can do away with this and simplify basically everything that touches images. The one potentially controversial change is to no longer separate internal and external images on disk. If we retain this separation, we have to keep `image_origin` around in a number of spots and it getting image paths on disk painful. So, I am have gotten rid of this organisation. Images are now all stored in `images`, regardless of their origin. As we improve the image management features, this change will hopefully become transparent.
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
from typing import Literal, Union
|
|
|
|
from pydantic import Field
|
|
|
|
from invokeai.app.models.image import ImageCategory, ImageField, ResourceOrigin
|
|
|
|
from .baseinvocation import BaseInvocation, InvocationContext, InvocationConfig
|
|
from .image import ImageOutput
|
|
|
|
|
|
class RestoreFaceInvocation(BaseInvocation):
|
|
"""Restores faces in an image."""
|
|
|
|
# fmt: off
|
|
type: Literal["restore_face"] = "restore_face"
|
|
|
|
# Inputs
|
|
image: Union[ImageField, None] = Field(description="The input image")
|
|
strength: float = Field(default=0.75, gt=0, le=1, description="The strength of the restoration" )
|
|
# fmt: on
|
|
|
|
# Schema customisation
|
|
class Config(InvocationConfig):
|
|
schema_extra = {
|
|
"ui": {
|
|
"tags": ["restoration", "image"],
|
|
},
|
|
}
|
|
|
|
def invoke(self, context: InvocationContext) -> ImageOutput:
|
|
image = context.services.images.get_pil_image(self.image.image_name)
|
|
results = context.services.restoration.upscale_and_reconstruct(
|
|
image_list=[[image, 0]],
|
|
upscale=None,
|
|
strength=self.strength, # GFPGAN strength
|
|
save_original=False,
|
|
image_callback=None,
|
|
)
|
|
|
|
# Results are image and seed, unwrap for now
|
|
# 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_name=image_dto.image_name),
|
|
width=image_dto.width,
|
|
height=image_dto.height,
|
|
)
|