From 12c31e8967868209ec15f6172ca7caaf56a0a661 Mon Sep 17 00:00:00 2001
From: VALADI K JAGANATHAN <info@orchidsasia.com>
Date: Thu, 30 Nov 2023 07:55:23 +0530
Subject: [PATCH] Update city96_latent_upscaler.py

added a line on noise mask
---
 py/city96_latent_upscaler.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/py/city96_latent_upscaler.py b/py/city96_latent_upscaler.py
index 1d0b9f4..f1656e9 100644
--- a/py/city96_latent_upscaler.py
+++ b/py/city96_latent_upscaler.py
@@ -1,4 +1,5 @@
 # https://github.com/city96/SD-Latent-Upscaler
+# check if any issues on weights and updates and local install for same
 import torch
 import torch.nn as nn
 from safetensors.torch import load_file
@@ -79,4 +80,9 @@ class LatentUpscaler:
 		lt = samples["samples"]
 		lt = model(lt)
 		del model
-		return ({"samples": lt},)
\ No newline at end of file
+		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},)
+