From 45b049669bf70414292104abe224aa956e02e86e Mon Sep 17 00:00:00 2001
From: "Lt.Dr.Data" <lt.dr.data@gmail.com>
Date: Thu, 27 Apr 2023 14:33:02 +0900
Subject: [PATCH] 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.
---
 server.py                         | 43 ++++++++++++++++++++++++++++---
 web/extensions/core/maskeditor.js | 26 +++++++++----------
 2 files changed, 52 insertions(+), 17 deletions(-)

diff --git a/server.py b/server.py
index 3732ee5bc..40c42df3b 100644
--- a/server.py
+++ b/server.py
@@ -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")
diff --git a/web/extensions/core/maskeditor.js b/web/extensions/core/maskeditor.js
index eb601720d..3ebac7bd2 100644
--- a/web/extensions/core/maskeditor.js
+++ b/web/extensions/core/maskeditor.js
@@ -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);