From 711d09a1070bc5aadc3cfeafda4152ffa162a612 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Tue, 25 Apr 2023 00:24:03 +1000
Subject: [PATCH] feat(nodes): add `get_uri` method to image storage

- gets the external URI of an image
---
 invokeai/app/api/routers/images.py     | 18 +++++++++--------
 invokeai/app/services/image_storage.py | 28 +++++++++++++++++++++++---
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/invokeai/app/api/routers/images.py b/invokeai/app/api/routers/images.py
index 22c07033cf..98656c8bda 100644
--- a/invokeai/app/api/routers/images.py
+++ b/invokeai/app/api/routers/images.py
@@ -104,17 +104,19 @@ async def upload_image(
 
     invokeai_metadata = ApiDependencies.invoker.services.metadata.get_metadata(img)
 
+    image_url = ApiDependencies.invoker.services.images.get_uri(
+        ImageType.UPLOAD, saved_image.image_name
+    )
+
+    thumbnail_url = ApiDependencies.invoker.services.images.get_uri(
+        ImageType.UPLOAD, saved_image.image_name, True
+    )
+
     res = ImageResponse(
         image_type=ImageType.UPLOAD,
         image_name=saved_image.image_name,
-        image_url=request.url_for(
-            "get_image", image_type=ImageType.UPLOAD.value, image_name=saved_image.image_name
-        ),
-        thumbnail_url=request.url_for(
-            "get_thumbnail",
-            thumbnail_type=ImageType.UPLOAD.value,
-            thumbnail_name=saved_image.thumbnail_name,
-        ),
+        image_url=image_url,
+        thumbnail_url=thumbnail_url,
         metadata=ImageResponseMetadata(
             created=saved_image.created,
             width=img.width,
diff --git a/invokeai/app/services/image_storage.py b/invokeai/app/services/image_storage.py
index cd2c8a13a3..bec4aca850 100644
--- a/invokeai/app/services/image_storage.py
+++ b/invokeai/app/services/image_storage.py
@@ -45,7 +45,15 @@ class ImageStorageBase(ABC):
     def get_path(
         self, image_type: ImageType, image_name: str, is_thumbnail: bool = False
     ) -> str:
-        """Gets the path to an image or its thumbnail."""
+        """Gets the internal path to an image or its thumbnail."""
+        pass
+
+    # TODO: make this a bit more flexible for e.g. cloud storage
+    @abstractmethod
+    def get_uri(
+        self, image_type: ImageType, image_name: str, is_thumbnail: bool = False
+    ) -> str:
+        """Gets the external URI to an image or its thumbnail."""
         pass
 
     # TODO: make this a bit more flexible for e.g. cloud storage
@@ -130,8 +138,8 @@ class DiskImageStorage(ImageStorageBase):
                     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",
+                    image_url=self.get_uri(image_type, filename),
+                    thumbnail_url=self.get_uri(image_type, filename, True),
                     # TODO: Creation of this object should happen elsewhere (?), just making it fit here so it works
                     metadata=ImageResponseMetadata(
                         created=int(os.path.getctime(path)),
@@ -180,6 +188,20 @@ class DiskImageStorage(ImageStorageBase):
 
         return path
 
+    def get_uri(
+        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:
+            thumbnail_basename = get_thumbnail_name(basename)
+            uri = f"api/v1/images/{image_type.value}/thumbnails/{thumbnail_basename}"
+        else:
+            uri = f"api/v1/images/{image_type.value}/{basename}"
+
+        return uri
+
     def validate_path(self, path: str) -> bool:
         try:
             os.stat(path)