From f0eb8206b1773bb45e8ed235cd848289a8ca5140 Mon Sep 17 00:00:00 2001
From: asagi4 <130366179+asagi4@users.noreply.github.com>
Date: Sun, 8 Dec 2024 18:33:54 +0200
Subject: [PATCH] This is the last thing for now, I swear

---
 README.md   |  4 ++--
 __init__.py | 39 ++++++++++++++++++++++++++-------------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index 22aa01b..729b2e2 100644
--- a/README.md
+++ b/README.md
@@ -12,10 +12,10 @@ The original pickle-format checkpoints are found at https://drive.google.com/dri
 ## Usage
 Use with custom sampling and pass in an initial noise from eg. `RandomNoise` and a cond (only the first prompt in the conditioning will be used if multiple exist).
 
-You can also run it on the CPU, though appears to change the output for some reason.
+You can also run it on the CPU, though that appears to change the output for some reason.
 
 ## Notes
-The model works with 128x128 latents, apparently. If you pass in other shaped latents, it will reshape the noise into a square before running the noise model, and then reshape the result back to the original resolution.
+The model works with 128x128 latents, apparently. If you pass in other shaped latents, it will reshape the noise into a square before running the noise model, and then reshape the result back to the original resolution. You can control how the reshape happens with the `reshape` and `method` parameters.
 
 If you get an error from the timm module when running this, update your timm package. It may be too old.
 
diff --git a/__init__.py b/__init__.py
index 65ef6ad..853cfac 100644
--- a/__init__.py
+++ b/__init__.py
@@ -202,6 +202,8 @@ class NPNetGoldenNoise:
     noise = None
     cond = None
     seed = None
+    method = "nearest-exact"
+    strategy = "resize"
 
     @classmethod
     def INPUT_TYPES(s):
@@ -217,7 +219,11 @@ class NPNetGoldenNoise:
                 "prompt": ("CONDITIONING",),
                 "model": (folder_paths.get_filename_list("npnet"),),
                 "device": (["cuda", "cpu"],),
-            }
+            },
+            "optional": {
+                "reshape": (["resize", "crop"],),
+                "method": (["nearest-exact", "bilinear", "area", "bicubic", "bislerp"],),
+            },
         }
 
     RETURN_TYPES = ("NOISE",)
@@ -225,30 +231,37 @@ class NPNetGoldenNoise:
 
     FUNCTION = "doit"
 
+    def reshape(self, noise, shape):
+        if shape[-1] == noise.shape[-1] and shape[-2] == noise.shape[-2]:
+            return noise
+        crop = "disabled" if self.strategy == "resize" else "center"
+        return common_upscale(noise, shape[-1], shape[-2], self.method, crop)
+
     def generate_noise(self, input_latent):
         self.seed = self.noise.seed
         orig_shape = input_latent["samples"].shape
-        if orig_shape[-2] != 128 or orig_shape[-1] != 128:
-            input_latent = input_latent.copy()
-            print("Latent must be 128x128 for the NPNet model to work; generating square noise and reshaping...")
-            input_latent["samples"] = common_upscale(input_latent["samples"], 128, 128, "nearest-exact", "disabled")
+        input_latent = input_latent.copy()
+        input_latent["samples"] = self.reshape(input_latent["samples"], (128, 128))
         init_noise = self.noise.generate_noise(input_latent).to(self.npnet.device)
         cond = self.cond[0].clone().to(self.npnet.device)
         if cond.shape[1] != 77:
-            print("NPNet can't handle conds >77 tokens, truncating...")
-            cond = cond[:, :77, :]
-        print("Applying NPNet to noise")
-        r = self.npnet(init_noise, cond).to("cpu")
-        if orig_shape[-2] != 128 or orig_shape[-1] != 128:
-            r = common_upscale(r, orig_shape[-1], orig_shape[-2], "nearest-exact", "disabled")
-        return r
+            print("NPNet can't handle conds >77 tokens, running the model individually on each piece")
 
-    def doit(self, noise, prompt, model, device):
+        r = init_noise
+        for i, cond in enumerate(torch.split(cond, 77, 1)):
+            print("Applying NPNet to chunk", i + 1)
+            r = self.npnet(r, cond)
+
+        return self.reshape(r.to("cpu"), orig_shape)
+
+    def doit(self, noise, prompt, model, device, reshape="resize", method="nearest-exact"):
         model_path = folder_paths.get_full_path("npnet", model)
         if self.npnet is None or self.npnet.pretrained_path != model_path:
             print("Loading NPNet from", model_path)
             self.npnet = NPNet(model_path, device=device)
         self.npnet.to(device)
+        self.method = method
+        self.strategy = reshape
         self.noise = noise
         self.cond = prompt[0]