From 790f399986ec5914969f5b1bbeacc549044cf63a Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Fri, 21 Apr 2023 19:13:26 +1000
Subject: [PATCH] feat(nodes): tidy images routes

---
 invokeai/app/api/routers/images.py     | 26 ++++++++++++++++----------
 invokeai/app/services/image_storage.py |  8 ++++----
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/invokeai/app/api/routers/images.py b/invokeai/app/api/routers/images.py
index da8208d981..2eb64cfd6b 100644
--- a/invokeai/app/api/routers/images.py
+++ b/invokeai/app/api/routers/images.py
@@ -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,
diff --git a/invokeai/app/services/image_storage.py b/invokeai/app/services/image_storage.py
index 1739810542..af40bf47ca 100644
--- a/invokeai/app/services/image_storage.py
+++ b/invokeai/app/services/image_storage.py
@@ -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):