### 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.
This commit is contained in:
Pupbani 2024-12-30 15:40:16 +09:00
parent 82ecb02c1e
commit 59844491b1

View File

@ -329,11 +329,46 @@ class PromptServer():
else: else:
return web.Response(status=400) 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") @routes.post("/upload/image")
async def upload_image(request): async def upload_image(request):
post = await request.post() post = await request.post()
return image_upload(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") @routes.post("/upload/mask")
async def upload_mask(request): async def upload_mask(request):
@ -378,6 +413,32 @@ class PromptServer():
return image_upload(post, image_save_function) 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") @routes.get("/view")
async def view_image(request): async def view_image(request):
if "filename" in request.rel_url.query: if "filename" in request.rel_url.query: