From 50ac3eb28dafedc5b100d8bf88560b02cf4e591e Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Fri, 21 Apr 2023 19:12:20 +1000
Subject: [PATCH] feat(nodes): add delete_image & delete_images routes

---
 invokeai/app/api/models/images.py      |  7 +++++
 invokeai/app/api/routers/images.py     | 38 ++++++++++++++++++++++++--
 invokeai/app/services/image_storage.py | 12 ++++----
 3 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/invokeai/app/api/models/images.py b/invokeai/app/api/models/images.py
index 53caca63f5..c0f2b152c7 100644
--- a/invokeai/app/api/models/images.py
+++ b/invokeai/app/api/models/images.py
@@ -32,3 +32,10 @@ class ProgressImage(BaseModel):
     width: int = Field(description="The effective width of the image in pixels")
     height: int = Field(description="The effective height of the image in pixels")
     dataURL: str = Field(description="The image data as a b64 data URL")
+
+
+class NonNullableImageField(BaseModel):
+    """Non-nullable ImageField, used for delete_images route"""
+
+    image_type: ImageType
+    image_name: str
diff --git a/invokeai/app/api/routers/images.py b/invokeai/app/api/routers/images.py
index 14a84a5dd4..da8208d981 100644
--- a/invokeai/app/api/routers/images.py
+++ b/invokeai/app/api/routers/images.py
@@ -6,12 +6,15 @@ import os
 from typing import Any
 import uuid
 
-from fastapi import HTTPException, Path, Query, Request, UploadFile
+from fastapi import Body, HTTPException, Path, Query, Request, UploadFile
 from fastapi.responses import FileResponse, Response
 from fastapi.routing import APIRouter
 from PIL import Image
-from invokeai.app.api.models.images import ImageResponse, ImageResponseMetadata
-from invokeai.app.services.metadata import InvokeAIMetadata
+from invokeai.app.api.models.images import (
+    ImageResponse,
+    ImageResponseMetadata,
+    NonNullableImageField,
+)
 from invokeai.app.services.item_storage import PaginatedResults
 
 from ...services.image_storage import ImageType
@@ -37,6 +40,18 @@ async def get_image(
         raise HTTPException(status_code=404)
 
 
+@images_router.delete("/{image_type}/{image_name}", operation_id="delete_image")
+async def delete_image(
+    image_type: ImageType = Path(description="The type of image to delete"),
+    image_name: str = Path(description="The name of the image to delete"),
+) -> None:
+    """Deletes an image and its thumbnail"""
+
+    ApiDependencies.invoker.services.images.delete(
+        image_type=image_type, image_name=image_name
+    )
+
+
 @images_router.get(
     "/{image_type}/thumbnails/{image_name}", operation_id="get_thumbnail"
 )
@@ -126,3 +141,20 @@ async def list_images(
     """Gets a list of images"""
     result = ApiDependencies.invoker.services.images.list(image_type, page, per_page)
     return result
+
+
+@images_router.delete(
+    "/",
+    operation_id="delete_images",
+)
+async def delete_images(
+    images: list[NonNullableImageField] = Body(
+        description="The list of images to delete"
+    ),
+) -> None:
+    """Deletes a list of images and their thumbnails"""
+
+    for i in images:
+        ApiDependencies.invoker.services.images.delete(
+            image_type=i.image_type, image_name=i.image_name
+        )
diff --git a/invokeai/app/services/image_storage.py b/invokeai/app/services/image_storage.py
index 335425fe66..1739810542 100644
--- a/invokeai/app/services/image_storage.py
+++ b/invokeai/app/services/image_storage.py
@@ -210,17 +210,19 @@ class DiskImageStorage(ImageStorageBase):
         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)
+        basename = os.path.basename(image_name)
+        image_path = self.get_path(image_type, basename)
+
         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)
+        thumbnail_name = f"{os.path.splitext(basename)[0]}.webp"
+        thumbnail_path = self.get_path(image_type, thumbnail_name, True)
 
+        if os.path.exists(image_path):
+            os.remove(thumbnail_path)
         if thumbnail_path in self.__cache:
             del self.__cache[thumbnail_path]