From a04a7cacd7169755a12a138b51abc177e91576a5 Mon Sep 17 00:00:00 2001
From: yoinked <pepe.dannyboy@gmail.com>
Date: Sun, 26 Nov 2023 15:47:07 -0800
Subject: [PATCH] added noise mask support

---
 comfy_latent_upscaler.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/comfy_latent_upscaler.py b/comfy_latent_upscaler.py
index 047542b..852aa62 100644
--- a/comfy_latent_upscaler.py
+++ b/comfy_latent_upscaler.py
@@ -90,6 +90,10 @@ class LatentUpscaler:
 		lt = samples["samples"]
 		lt = model(lt)
 		del model
+		if "noise_mask" in samples.keys():
+			# expand the noise mask to the same shape as the latent
+			mask = torch.nn.functional.interpolate(samples['noise_mask'], scale_factor=float(scale_factor), mode='bicubic') 
+			return ({"samples": lt, "noise_mask": mask},)
 		return ({"samples": lt},)
 
 NODE_CLASS_MAPPINGS = {