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 <lt.dr.data@gmail.com>

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)

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");
 

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);
-						}
 					});
 				}
 			}