Skip to content

Commit

Permalink
optimize: support preview mode for mask editor. (comfyanonymous#755)
Browse files Browse the repository at this point in the history 
* 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 <[email protected]>
  • Loading branch information
@ltdrdata
ltdrdata and Lt.Dr.Data authored Jun 24, 2023
1 parent 0567694 commit c9f5d5b
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 20 deletions.
44 changes: 34 additions & 10 deletions server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -186,18 +186,43 @@ async def upload_image(request):
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'])

# 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

# alpha copy
new_alpha = mask_pil.getchannel('A')
original_pil.putalpha(new_alpha)
original_pil.save(filepath, compress_level=4)
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)

Expand Down Expand Up @@ -231,17 +256,16 @@ async def view_image(request):
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
if preview_info[-1].isdigit():
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)
Expand Down
15 changes: 12 additions & 3 deletions web/extensions/core/maskeditor.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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");

Expand Down
15 changes: 8 additions & 7 deletions web/scripts/app.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -159,25 +159,26 @@ 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 {
prop.value = value;
prop.callback(value);
}
}
else if (prop && prop.type != 'button') {
prop.value = value;
prop.callback(value);
}
});
}
}
Expand Down

0 comments on commit c9f5d5b

Please sign in to comment.