From b4f2843a8b4c374188f9f96d11c43ea5fa96727c Mon Sep 17 00:00:00 2001
From: chflame163 <130118553+chflame163@users.noreply.github.com>
Date: Thu, 18 Jan 2024 00:02:06 +0800
Subject: [PATCH] fix mask output bug

---
 py/comfy_wordcloud.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/py/comfy_wordcloud.py b/py/comfy_wordcloud.py
index 4be10c8..3bfb2a3 100644
--- a/py/comfy_wordcloud.py
+++ b/py/comfy_wordcloud.py
@@ -58,6 +58,10 @@ def tensor2pil(image):
 def pil2tensor(image):
     return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
 
+def getRGBAmask(image):
+    ret_mask = torch.tensor([pil2tensor(image)[0, :, :, 3].tolist()])
+    return ret_mask
+
 def img_whitebackground(image):
     if image.mode != 'RGBA':
         image = image.convert('RGBA')
@@ -65,7 +69,6 @@ def img_whitebackground(image):
     height = image.height
     img_new = Image.new('RGB', size=(width, height), color=(255, 255, 255))
     img_new.paste(image, (0, 0), mask=image)
-
     return img_new
 
 class ComfyWordCloud:
@@ -201,7 +204,7 @@ class ComfyWordCloud:
             wc.recolor(color_func=image_colors)
 
         ret_image = wc.to_image().convert('RGBA')
-        ret_mask = pil2tensor(ret_image)[0, :, :, 3]
+        ret_mask = getRGBAmask(ret_image)
 
         return (pil2tensor(ret_image), ret_mask,)