Add support for channel-specific image data retrieval in /view API to fix alpha mask loading issue

When loading an image with an alpha mask in JavaScript canvas, there is an issue where the alpha and RGB channels are premultiplied. To avoid reliance on JavaScript canvas, I added support for channel-specific image data retrieval in the "/view" API. This allows us to retrieve data for each channel separately and fix the alpha mask loading issue. The changes have been committed to the repository.
This commit is contained in:
Lt.Dr.Data 2023-04-27 14:33:02 +09:00
parent a13fee6882
commit 45b049669b
2 changed files with 52 additions and 17 deletions

View File

@ -8,7 +8,7 @@ import uuid
import json
import glob
from PIL import Image
import numpy as np
from io import BytesIO
try:
import aiohttp
@ -216,8 +216,45 @@ class PromptServer():
file = os.path.join(output_dir, filename)
if os.path.isfile(file):
return web.FileResponse(file, headers={"Content-Disposition": f"filename=\"{filename}\""})
if 'channel' not in request.rel_url.query:
channel = 'rgba'
else:
channel = request.rel_url.query["channel"]
if channel == 'rgb':
with Image.open(file) as img:
if img.mode == "RGBA":
r, g, b, a = img.split()
new_img = Image.merge('RGB', (r, g, b))
else:
new_img = img.convert("RGB")
buffer = BytesIO()
new_img.save(buffer, format='PNG')
buffer.seek(0)
return web.Response(body=buffer.read(), content_type='image/png',
headers={"Content-Disposition": f"filename=\"{filename}\""})
elif channel == 'a':
with Image.open(file) as img:
if img.mode == "RGBA":
_, _, _, a = img.split()
else:
a = Image.new('L', img.size, 255)
# alpha img
alpha_img = Image.new('RGBA', img.size)
alpha_img.putalpha(a)
alpha_buffer = BytesIO()
alpha_img.save(alpha_buffer, format='PNG')
alpha_buffer.seek(0)
return web.Response(body=alpha_buffer.read(), content_type='image/png',
headers={"Content-Disposition": f"filename=\"{filename}\""})
else:
return web.FileResponse(file, headers={"Content-Disposition": f"filename=\"{filename}\""})
return web.Response(status=404)
@routes.get("/prompt")

View File

@ -45,7 +45,7 @@ async function uploadMask(filepath, formData) {
ComfyApp.clipspace.images = [filepath];
}
function removeRGB(image, backupCanvas, backupCtx, maskCtx) {
function prepareRGB(image, backupCanvas, backupCtx) {
// paste mask data into alpha channel
backupCtx.drawImage(image, 0, 0, backupCanvas.width, backupCanvas.height);
const backupData = backupCtx.getImageData(0, 0, backupCanvas.width, backupCanvas.height);
@ -218,19 +218,11 @@ class MaskEditorDialog extends ComfyDialog {
this.setlayout(imgCanvas, maskCanvas);
// prepare content
this.maskCanvas = maskCanvas;
this.backupCanvas = backupCanvas;
this.maskCtx = maskCanvas.getContext('2d');
this.backupCtx = backupCanvas.getContext('2d');
// separate original_imgs and imgs
if(ComfyApp.clipspace.imgs[0] === ComfyApp.clipspace.original_imgs[0]) {
var copiedImage = new Image();
copiedImage.src = ComfyApp.clipspace.original_imgs[0].src;
ComfyApp.clipspace.imgs = [copiedImage];
}
this.setImages(imgCanvas, backupCanvas);
this.setEventHandler(maskCanvas);
}
@ -280,19 +272,25 @@ class MaskEditorDialog extends ComfyDialog {
backupCanvas.width = touched_image.width;
backupCanvas.height = touched_image.height;
removeRGB(touched_image, backupCanvas, backupCtx, maskCtx);
prepareRGB(touched_image, backupCanvas, backupCtx);
};
touched_image.src = ComfyApp.clipspace.imgs[0].src;
const alpha_url = new URL(ComfyApp.clipspace.imgs[0].src)
alpha_url.searchParams.delete('channel');
alpha_url.searchParams.set('channel', 'a');
touched_image.src = alpha_url;
// original image load
orig_image.onload = function() {
window.dispatchEvent(new Event('resize'));
};
orig_image.src = ComfyApp.clipspace.original_imgs[0].src;
const rgb_url = new URL(ComfyApp.clipspace.imgs[0].src);
rgb_url.searchParams.delete('channel');
rgb_url.searchParams.set('channel', 'rgb');
orig_image.src = rgb_url;
this.image = orig_image;
}
}g
setEventHandler(maskCanvas) {
@ -523,7 +521,7 @@ class MaskEditorDialog extends ComfyDialog {
const dataURL = this.backupCanvas.toDataURL();
const blob = dataURLToBlob(dataURL);
const original_blob = loadedImageToBlob(ComfyApp.clipspace.original_imgs[0]);
const original_blob = loadedImageToBlob(this.image);
formData.append('image', blob, filename);
formData.append('original_image', original_blob);