From e7654ba3515d76f1633e3309f2421139265e66c7 Mon Sep 17 00:00:00 2001
From: chflame163 <130118553+chflame163@users.noreply.github.com>
Date: Mon, 15 Jan 2024 11:52:22 +0800
Subject: [PATCH] fix error of 'ImageColorGenerator is smaller than the canvas'

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

diff --git a/py/comfy_wordcloud.py b/py/comfy_wordcloud.py
index b0b0b49..9d30973 100644
--- a/py/comfy_wordcloud.py
+++ b/py/comfy_wordcloud.py
@@ -177,8 +177,14 @@ class ComfyWordCloud:
             random_state = None
 
         mask = None
+        image_width = width
+        image_height = height
         if not mask_image == None:
-            mask = np.array(img_whitebackground(tensor2pil(mask_image)))
+            p_mask = tensor2pil(mask_image)
+            mask = np.array(img_whitebackground(p_mask))
+            image_width = p_mask.width
+            image_height = p_mask.height
+
 
         # set wordcloud parameters
         wc = WordCloud(width=width, height=height, scale=scale, margin=margin,
@@ -195,7 +201,9 @@ class ComfyWordCloud:
 
         # generate recolor
         if not color_ref_image == None:
-            image_colors = ImageColorGenerator(np.array(tensor2pil(color_ref_image)))
+            p_color_ref_image = tensor2pil(color_ref_image)
+            p_color_ref_image = p_color_ref_image.resize((image_width, image_height))
+            image_colors = ImageColorGenerator(np.array(p_color_ref_image))
             wc.recolor(color_func=image_colors)
 
         return (pil2tensor(wc.to_image()),)