From 59844491b1733b16a11c23e525514902b6cba89e Mon Sep 17 00:00:00 2001
From: Pupbani <rhdk5148@gmail.com>
Date: Mon, 30 Dec 2024 15:40:16 +0900
Subject: [PATCH] ### Title Add image upload and delete endpoints in server.py

### Description
This pull request introduces new endpoints in `server.py` for uploading images in byte stream format and for deleting files.

### Changes Made
- Added an endpoint for uploading images: `POST /upload/image_stream`
- Added an endpoint for deleting files: `DELETE /delete-file`

### Notes
Please review the changes and let me know if any modifications are needed.
---
 server.py | 61 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)

diff --git a/server.py b/server.py
index 2dc53b9d..b7cc0acf 100644
--- a/server.py
+++ b/server.py
@@ -329,11 +329,46 @@ class PromptServer():
             else:
                 return web.Response(status=400)
 
+        def str2pil(img_str:str)->Image.Image:
+            """
+            Convert Image Byte Stream to PIL Image
+            """
+            import io,base64
+
+            img_data = base64.b64decode(img_str)
+            img_io = io.BytesIO(img_data)
+            img = Image.open(img_io)
+            return img
+
         @routes.post("/upload/image")
         async def upload_image(request):
             post = await request.post()
             return image_upload(post)
 
+        @routes.post("/upload/image_stream")
+        async def upload_image(request):
+            post = await request.post()
+            
+            # Get Image Byte Stream
+            image = str2pil(post.get("img_str"))
+            
+            file_name = post.get("file_name")
+            if not(file_name.endswith(".png") or file_name.endswith(".jpg")):
+                return web.Response(text="The file name must end in .jpg or .png.", status=400)
+            
+            # PATH
+            if post.get("subfolder","") =="":
+                UPLOAD_PATH = os.path.join(os.getcwd(),post.get("type"),file_name)
+            else:
+                sub = post.get("subfolder")
+                sub_dir_path = os.path.join(os.getcwd(),post.get("type"),post.get("subfolder"))
+                os.makedirs(sub_dir_path,exist_ok=True)
+                UPLOAD_PATH = os.path.join(sub_dir_path,file_name)
+
+            image.save(UPLOAD_PATH)
+
+            return web.Response(status=200,text=f"Success Save Image PATH :{UPLOAD_PATH}")
+
 
         @routes.post("/upload/mask")
         async def upload_mask(request):
@@ -378,6 +413,32 @@ class PromptServer():
 
             return image_upload(post, image_save_function)
 
+        @routes.delete("/delete-file")
+        async def delete_file(request):
+            type = request.query.get("type")
+            if not type:
+                return web.Response(text="Need Type",status=400)
+        
+            filename = request.query.get("file_name")
+            if not filename:
+                return web.Response(text="Need Filename",status=400)
+        
+            sub = request.query.get("subfolder","")
+        
+            if sub == "":
+                file_path = os.path.join(os.getcwd(),type,filename)
+            else:
+                file_path = os.path.join(os.getcwd(),type,sub,filename)
+            
+            if not os.path.isfile(file_path):
+                return web.Response(text="Not Found File", status=404)
+            
+            try:
+                os.remove(file_path)
+                return web.Response(text=f"Success Remove File : {file_path[len(os.getcwd())+1:]}",status=200)
+            except Exception as e:
+                return web.Response(text=f"Failed Remove File\n{str(e)}",status=500)
+
         @routes.get("/view")
         async def view_image(request):
             if "filename" in request.rel_url.query: