Skip to content

Commit

Permalink
feat(nodes): add thumbnail generation to DiskImageStorage
Browse files Browse the repository at this point in the history
  • Loading branch information
psychedelicious committed Apr 4, 2023
1 parent 18cddd7 commit aed9ece
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 1 deletion.
10 changes: 10 additions & 0 deletions invokeai/app/api/routers/images.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,16 @@ async def get_image(
filename = ApiDependencies.invoker.services.images.get_path(image_type, image_name)
return FileResponse(filename)

@images_router.get("/{image_type}/thumbnails/{image_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"),
):
"""Gets a thumbnail"""
# TODO: This is not really secure at all. At least make sure only output results are served
filename = ApiDependencies.invoker.services.images.get_path(image_type, 'thumbnails/' + image_name)
return FileResponse(filename)


@images_router.post(
"/uploads/",
Expand Down
10 changes: 9 additions & 1 deletion invokeai/app/services/image_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from typing import Dict

from PIL.Image import Image
from invokeai.app.util.save_thumbnail import save_thumbnail

from invokeai.backend.image_util import PngWriter

Expand Down Expand Up @@ -66,6 +67,9 @@ def __init__(self, output_folder: str):
Path(os.path.join(output_folder, image_type)).mkdir(
parents=True, exist_ok=True
)
Path(os.path.join(output_folder, image_type, "thumbnails")).mkdir(
parents=True, exist_ok=True
)

def get(self, image_type: ImageType, image_name: str) -> Image:
image_path = self.get_path(image_type, image_name)
Expand All @@ -87,7 +91,11 @@ def save(self, image_type: ImageType, image_name: str, image: Image) -> None:
self.__pngWriter.save_image_and_prompt_to_png(
image, "", image_subpath, None
) # TODO: just pass full path to png writer

save_thumbnail(
image=image,
filename=image_name,
path=os.path.join(self.__output_folder, image_type, "thumbnails"),
)
image_path = self.get_path(image_type, image_name)
self.__set_cache(image_path, image)

Expand Down
25 changes: 25 additions & 0 deletions invokeai/app/util/save_thumbnail.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
import os
from PIL import Image


def save_thumbnail(
image: Image.Image,
filename: str,
path: str,
size: int = 256,
) -> str:
"""
Saves a thumbnail of an image, returning its path.
"""
base_filename = os.path.splitext(filename)[0]
thumbnail_path = os.path.join(path, base_filename + ".webp")

if os.path.exists(thumbnail_path):
return thumbnail_path

image_copy = image.copy()
image_copy.thumbnail(size=(size, size))

image_copy.save(thumbnail_path, "WEBP")

return thumbnail_path

0 comments on commit aed9ece

Please sign in to comment.