From c9f5d5b2e12a3b0bd3232927a617ef646418dee5 Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" <128333288+ltdrdata@users.noreply.github.com> Date: Sat, 24 Jun 2023 16:45:41 +0900 Subject: [PATCH] optimize: support preview mode for mask editor. (#755) * support preview mode for mask editor. * use original file reference instead of loaded frontend blob bugfix: * prevent file open dialog when save to load image * bugfix: cannot clear previous mask painted image's alpha * bugfix * bugfix --------- Co-authored-by: Lt.Dr.Data --- server.py | 44 ++++++++++++++++++++++++------- web/extensions/core/maskeditor.js | 15 ++++++++--- web/scripts/app.js | 15 ++++++----- 3 files changed, 54 insertions(+), 20 deletions(-) diff --git a/server.py b/server.py index f385cefb..7b4fcac3 100644 --- a/server.py +++ b/server.py @@ -64,7 +64,7 @@ class PromptServer(): def __init__(self, loop): PromptServer.instance = self - mimetypes.init(); + mimetypes.init() mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8' self.prompt_queue = None self.loop = loop @@ -186,18 +186,43 @@ class PromptServer(): post = await request.post() return image_upload(post) + @routes.post("/upload/mask") async def upload_mask(request): post = await request.post() def image_save_function(image, post, filepath): - original_pil = Image.open(post.get("original_image").file).convert('RGBA') - mask_pil = Image.open(image.file).convert('RGBA') + original_ref = json.loads(post.get("original_ref")) + filename, output_dir = folder_paths.annotated_filepath(original_ref['filename']) - # alpha copy - new_alpha = mask_pil.getchannel('A') - original_pil.putalpha(new_alpha) - original_pil.save(filepath, compress_level=4) + # validation for security: prevent accessing arbitrary path + if filename[0] == '/' or '..' in filename: + return web.Response(status=400) + + if output_dir is None: + type = original_ref.get("type", "output") + output_dir = folder_paths.get_directory_by_type(type) + + if output_dir is None: + return web.Response(status=400) + + if original_ref.get("subfolder", "") != "": + full_output_dir = os.path.join(output_dir, original_ref["subfolder"]) + if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: + return web.Response(status=403) + output_dir = full_output_dir + + file = os.path.join(output_dir, filename) + + if os.path.isfile(file): + with Image.open(file) as original_pil: + original_pil = original_pil.convert('RGBA') + mask_pil = Image.open(image.file).convert('RGBA') + + # alpha copy + new_alpha = mask_pil.getchannel('A') + original_pil.putalpha(new_alpha) + original_pil.save(filepath, compress_level=4) return image_upload(post, image_save_function) @@ -231,9 +256,8 @@ class PromptServer(): if 'preview' in request.rel_url.query: with Image.open(file) as img: preview_info = request.rel_url.query['preview'].split(';') - image_format = preview_info[0] - if image_format not in ['webp', 'jpeg']: + if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''): image_format = 'webp' quality = 90 @@ -241,7 +265,7 @@ class PromptServer(): quality = int(preview_info[-1]) buffer = BytesIO() - if image_format in ['jpeg']: + if image_format in ['jpeg'] or request.rel_url.query.get('channel', '') == 'rgb': img = img.convert("RGB") img.save(buffer, format=image_format, quality=quality) buffer.seek(0) diff --git a/web/extensions/core/maskeditor.js b/web/extensions/core/maskeditor.js index 764164d5..503c45f0 100644 --- a/web/extensions/core/maskeditor.js +++ b/web/extensions/core/maskeditor.js @@ -346,7 +346,6 @@ class MaskEditorDialog extends ComfyDialog { const rgb_url = new URL(ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src); rgb_url.searchParams.delete('channel'); - rgb_url.searchParams.delete('preview'); rgb_url.searchParams.set('channel', 'rgb'); orig_image.src = rgb_url; this.image = orig_image; @@ -618,10 +617,20 @@ class MaskEditorDialog extends ComfyDialog { const dataURL = this.backupCanvas.toDataURL(); const blob = dataURLToBlob(dataURL); - const original_blob = loadedImageToBlob(this.image); + let original_url = new URL(this.image.src); + + const original_ref = { filename: original_url.searchParams.get('filename') }; + + let original_subfolder = original_url.searchParams.get("subfolder"); + if(original_subfolder) + original_ref.subfolder = original_subfolder; + + let original_type = original_url.searchParams.get("type"); + if(original_type) + original_ref.type = original_type; formData.append('image', blob, filename); - formData.append('original_image', original_blob); + formData.append('original_ref', JSON.stringify(original_ref)); formData.append('type', "input"); formData.append('subfolder', "clipspace"); diff --git a/web/scripts/app.js b/web/scripts/app.js index 385a5457..4e83c40a 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -159,14 +159,19 @@ export class ComfyApp { const clip_image = ComfyApp.clipspace.images[ComfyApp.clipspace['selectedIndex']]; const index = node.widgets.findIndex(obj => obj.name === 'image'); if(index >= 0) { - node.widgets[index].value = clip_image; + if(node.widgets[index].type != 'image' && typeof node.widgets[index].value == "string" && clip_image.filename) { + node.widgets[index].value = (clip_image.subfolder?clip_image.subfolder+'/':'') + clip_image.filename + (clip_image.type?` [${clip_image.type}]`:''); + } + else { + node.widgets[index].value = clip_image; + } } } if(ComfyApp.clipspace.widgets) { ComfyApp.clipspace.widgets.forEach(({ type, name, value }) => { const prop = Object.values(node.widgets).find(obj => obj.type === type && obj.name === name); - if (prop && prop.type != 'image') { - if(typeof prop.value == "string" && value.filename) { + if (prop && prop.type != 'button') { + if(prop.type != 'image' && typeof prop.value == "string" && value.filename) { prop.value = (value.subfolder?value.subfolder+'/':'') + value.filename + (value.type?` [${value.type}]`:''); } else { @@ -174,10 +179,6 @@ export class ComfyApp { prop.callback(value); } } - else if (prop && prop.type != 'button') { - prop.value = value; - prop.callback(value); - } }); } }