InvokeAI/invokeai/app/services/image_storage.py

# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)

import os
from glob import glob
from abc import ABC, abstractmethod
from pathlib import Path
from queue import Queue
from typing import Dict, List, Tuple

from PIL.Image import Image
import PIL.Image as PILImage
from invokeai.app.api.models.images import ImageResponse, ImageResponseMetadata
from invokeai.app.models.image import ImageType
from invokeai.app.services.metadata import (
    InvokeAIMetadata,
    MetadataServiceBase,
    build_invokeai_metadata_pnginfo,
)
from invokeai.app.services.item_storage import PaginatedResults
from invokeai.app.util.misc import get_timestamp
from invokeai.app.util.thumbnails import get_thumbnail_name, make_thumbnail


class ImageStorageBase(ABC):
    """Responsible for storing and retrieving images."""

    @abstractmethod
    def get(self, image_type: ImageType, image_name: str) -> Image:
        """Retrieves an image as PIL Image."""
        pass

    @abstractmethod
    def list(
        self, image_type: ImageType, page: int = 0, per_page: int = 10
    ) -> PaginatedResults[ImageResponse]:
        """Gets a paginated list of images."""
        pass

    # TODO: make this a bit more flexible for e.g. cloud storage
    @abstractmethod
    def get_path(
        self, image_type: ImageType, image_name: str, is_thumbnail: bool = False
    ) -> str:
        """Gets the path to an image or its thumbnail."""
        pass

    # TODO: make this a bit more flexible for e.g. cloud storage
    @abstractmethod
    def validate_path(self, path: str) -> bool:
        """Validates an image path."""
        pass

    @abstractmethod
    def save(
        self,
        image_type: ImageType,
        image_name: str,
        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."""
        pass

    @abstractmethod
    def delete(self, image_type: ImageType, image_name: str) -> None:
        """Deletes an image and its thumbnail (if one exists)."""
        pass

    def create_name(self, context_id: str, node_id: str) -> str:
        """Creates a unique contextual image filename."""
        return f"{context_id}_{node_id}_{str(get_timestamp())}.png"


class DiskImageStorage(ImageStorageBase):
    """Stores images on disk"""

    __output_folder: str
    __cache_ids: Queue  # TODO: this is an incredibly naive cache
    __cache: Dict[str, Image]
    __max_cache_size: int
    __metadata_service: MetadataServiceBase

    def __init__(self, output_folder: str, metadata_service: MetadataServiceBase):
        self.__output_folder = output_folder
        self.__cache = dict()
        self.__cache_ids = Queue()
        self.__max_cache_size = 10  # TODO: get this from config
        self.__metadata_service = metadata_service

        Path(output_folder).mkdir(parents=True, exist_ok=True)

        # TODO: don't hard-code. get/save/delete should maybe take subpath?
        for image_type in ImageType:
            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 list(
        self, image_type: ImageType, page: int = 0, per_page: int = 10
    ) -> PaginatedResults[ImageResponse]:
        dir_path = os.path.join(self.__output_folder, image_type)
        image_paths = glob(f"{dir_path}/*.png")
        count = len(image_paths)

        sorted_image_paths = sorted(
            glob(f"{dir_path}/*.png"), key=os.path.getctime, reverse=True
        )

        page_of_image_paths = sorted_image_paths[
            page * per_page : (page + 1) * per_page
        ]

        page_of_images: List[ImageResponse] = []

        for path in page_of_image_paths:
            filename = os.path.basename(path)
            img = PILImage.open(path)

            invokeai_metadata = self.__metadata_service.get_metadata(img)

            page_of_images.append(
                ImageResponse(
                    image_type=image_type.value,
                    image_name=filename,
                    # TODO: DiskImageStorage should not be building URLs...?
                    image_url=f"api/v1/images/{image_type.value}/{filename}",
                    thumbnail_url=f"api/v1/images/{image_type.value}/thumbnails/{os.path.splitext(filename)[0]}.webp",
                    # TODO: Creation of this object should happen elsewhere (?), just making it fit here so it works
                    metadata=ImageResponseMetadata(
                        created=int(os.path.getctime(path)),
                        width=img.width,
                        height=img.height,
                        invokeai=invokeai_metadata,
                    ),
                )
            )

        page_count_trunc = int(count / per_page)
        page_count_mod = count % per_page
        page_count = page_count_trunc if page_count_mod == 0 else page_count_trunc + 1

        return PaginatedResults[ImageResponse](
            items=page_of_images,
            page=page,
            pages=page_count,
            per_page=per_page,
            total=count,
        )

    def get(self, image_type: ImageType, image_name: str) -> Image:
        image_path = self.get_path(image_type, image_name)
        cache_item = self.__get_cache(image_path)
        if cache_item:
            return cache_item

        image = PILImage.open(image_path)
        self.__set_cache(image_path, image)
        return image

    # TODO: make this a bit more flexible for e.g. cloud storage
    def get_path(
        self, image_type: ImageType, image_name: str, is_thumbnail: bool = False
    ) -> str:
        # strip out any relative path shenanigans
        basename = os.path.basename(image_name)

        if is_thumbnail:
            path = os.path.join(
                self.__output_folder, image_type, "thumbnails", basename
            )
        else:
            path = os.path.join(self.__output_folder, image_type, basename)

        return path

    def validate_path(self, path: str) -> bool:
        try:
            os.stat(path)
            return True
        except Exception:
            return False

    def save(
        self,
        image_type: ImageType,
        image_name: str,
        image: Image,
        metadata: InvokeAIMetadata | None = None,
    ) -> Tuple[str, str, int]:
        image_path = self.get_path(image_type, image_name)

        # TODO: Reading the image and then saving it strips the metadata...
        if metadata:
            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

        thumbnail_name = get_thumbnail_name(image_name)
        thumbnail_path = self.get_path(image_type, thumbnail_name, is_thumbnail=True)
        thumbnail_image = make_thumbnail(image)
        thumbnail_image.save(thumbnail_path)

        self.__set_cache(image_path, image)
        self.__set_cache(thumbnail_path, thumbnail_image)

        return (image_path, thumbnail_path, int(os.path.getctime(image_path)))

    def delete(self, image_type: ImageType, image_name: str) -> None:
        image_path = self.get_path(image_type, image_name)
        thumbnail_path = self.get_path(image_type, image_name, True)
        if os.path.exists(image_path):
            os.remove(image_path)

        if image_path in self.__cache:
            del self.__cache[image_path]

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

        if thumbnail_path in self.__cache:
            del self.__cache[thumbnail_path]

    def __get_cache(self, image_name: str) -> Image:
        return None if image_name not in self.__cache else self.__cache[image_name]

    def __set_cache(self, image_name: str, image: Image):
        if not image_name in self.__cache:
            self.__cache[image_name] = image
            self.__cache_ids.put(
                image_name
            )  # TODO: this should refresh position for LRU cache
            if len(self.__cache) > self.__max_cache_size:
                cache_id = self.__cache_ids.get()
                del self.__cache[cache_id]