feat(nodes): tidy images routes

This commit is contained in:
psychedelicious 2023-04-21 19:13:26 +10:00
parent 73cdd36594
commit 790f399986
2 changed files with 20 additions and 14 deletions

View File

@ -27,8 +27,8 @@ images_router = APIRouter(prefix="/v1/images", tags=["images"])
async def get_image(
image_type: ImageType = Path(description="The type of image to get"),
image_name: str = Path(description="The name of the image to get"),
) -> FileResponse | Response:
"""Gets a result"""
) -> FileResponse:
"""Gets an image"""
path = ApiDependencies.invoker.services.images.get_path(
image_type=image_type, image_name=image_name
@ -53,16 +53,16 @@ async def delete_image(
@images_router.get(
"/{image_type}/thumbnails/{image_name}", operation_id="get_thumbnail"
"/{thumbnail_type}/thumbnails/{thumbnail_name}", operation_id="get_thumbnail"
)
async def get_thumbnail(
image_type: ImageType = Path(description="The type of image to get"),
image_name: str = Path(description="The name of the image to get"),
thumbnail_type: ImageType = Path(description="The type of thumbnail to get"),
thumbnail_name: str = Path(description="The name of the thumbnail to get"),
) -> FileResponse | Response:
"""Gets a thumbnail"""
path = ApiDependencies.invoker.services.images.get_path(
image_type=image_type, image_name=image_name, is_thumbnail=True
image_type=thumbnail_type, image_name=thumbnail_name, is_thumbnail=True
)
if ApiDependencies.invoker.services.images.validate_path(path):
@ -99,7 +99,7 @@ async def upload_image(
filename = f"{uuid.uuid4()}_{str(int(datetime.now(timezone.utc).timestamp()))}.png"
(image_path, thumbnail_path, ctime) = ApiDependencies.invoker.services.images.save(
(image_name, thumbnail_name, ctime) = ApiDependencies.invoker.services.images.save(
ImageType.UPLOAD, filename, img
)
@ -107,9 +107,15 @@ async def upload_image(
res = ImageResponse(
image_type=ImageType.UPLOAD,
image_name=filename,
image_url=f"api/v1/images/{ImageType.UPLOAD.value}/{filename}",
thumbnail_url=f"api/v1/images/{ImageType.UPLOAD.value}/thumbnails/{os.path.splitext(filename)[0]}.webp",
image_name=image_name,
image_url=request.url_for(
"get_image", image_type=ImageType.UPLOAD.value, image_name=image_name
),
thumbnail_url=request.url_for(
"get_thumbnail",
thumbnail_type=ImageType.UPLOAD.value,
thumbnail_name=thumbnail_name,
),
metadata=ImageResponseMetadata(
created=ctime,
width=img.width,

View File

@ -58,7 +58,7 @@ class ImageStorageBase(ABC):
image: Image,
metadata: InvokeAIMetadata | None = None,
) -> Tuple[str, str, int]:
"""Saves an image and a 256x256 WEBP thumbnail. Returns a tuple of the image path, thumbnail path, and created timestamp."""
"""Saves an image and a 256x256 WEBP thumbnail. Returns a tuple of the image name, thumbnail name, and created timestamp."""
pass
@abstractmethod
@ -197,7 +197,7 @@ class DiskImageStorage(ImageStorageBase):
pnginfo = build_invokeai_metadata_pnginfo(metadata=metadata)
image.save(image_path, "PNG", pnginfo=pnginfo)
else:
image.save(image_path) # this saved image has an empty info
image.save(image_path) # this saved image has an empty info
thumbnail_name = get_thumbnail_name(image_name)
thumbnail_path = self.get_path(image_type, thumbnail_name, is_thumbnail=True)
@ -207,7 +207,7 @@ class DiskImageStorage(ImageStorageBase):
self.__set_cache(image_path, image)
self.__set_cache(thumbnail_path, thumbnail_image)
return (image_path, thumbnail_path, int(os.path.getctime(image_path)))
return (image_name, thumbnail_name, int(os.path.getctime(image_path)))
def delete(self, image_type: ImageType, image_name: str) -> None:
basename = os.path.basename(image_name)
@ -226,7 +226,7 @@ class DiskImageStorage(ImageStorageBase):
if thumbnail_path in self.__cache:
del self.__cache[thumbnail_path]
def __get_cache(self, image_name: str) -> Image:
def __get_cache(self, image_name: str) -> Image | None:
return None if image_name not in self.__cache else self.__cache[image_name]
def __set_cache(self, image_name: str, image: Image):