From 173230f800ccf79e4e270139dcf49da43293e57d Mon Sep 17 00:00:00 2001
From: Austin Mroz <AustinMroz@utexas.edu>
Date: Wed, 29 May 2024 20:33:45 -0500
Subject: [PATCH] Debug scaffolding, improved interpolation

When interpolating the subrender location, the mode is now explicitly
set so it does not default to nearest.

Added code to terminate sampling on after first subrender to assist in
visualize the correctness of the operation.

Added code to allow more convenient setting of the weighting applied to
the subrender.
---
 nodes.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/nodes.py b/nodes.py
index 6b40a27..5c84abe 100644
--- a/nodes.py
+++ b/nodes.py
@@ -50,6 +50,7 @@ def sample_dynamic(model, x, sigmas, extra_args=None, callback=None, disable=Non
     ch = h-h//4+1
     kernel = torch.ones((1,1,h//4,w//4), dtype=x.dtype, device=x.device)
     prev_denoised = None
+    early_terminate = False
     for i in trange(len(sigmas) - 1, disable=disable):
         gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
         sigma_hat = sigmas[i] * (gamma + 1)
@@ -80,19 +81,24 @@ def sample_dynamic(model, x, sigmas, extra_args=None, callback=None, disable=Non
                     sublocs.append(subx)
                     subinds.append(j)
             if len(sublocs) > 0:
+                early_terminate = True
                 subx = torch.stack(sublocs)
                 #subx = subx.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3)
-                subx = interpolate(subx, size=(h//2,w//2))
+                subx = interpolate(subx, size=(h//2,w//2), mode='bicubic')
                 subx += torch.randn_like(subx)*(1-math.sqrt(2))*sigma_hat*s_in
                 print("doing subrender")
                 subdn = model(subx, sigma_hat*s_in*math.sqrt(2), **extra_args)
                 #downscale
-                subdn = interpolate(subdn, size=(h//4,w//4))
+                subdn = interpolate(subdn, size=(h//4,w//4), mode='bicubic')
                 #TODO: Does added noise need to be filtered out?
                 for j,dn in zip(subinds,subdn):
                     subx = denoised[j,:,focal[j][0]:focal[j][0]+h//4,focal[j][1]:focal[j][1]+w//4]
-                    subx += 2*dn
-                    subx /= 3
+                    sub_weight = -1
+                    if sub_weight == -1:
+                        subx[:] = dn
+                    else:
+                        subx[:] += sub_weight*dn
+                        subx /= sub_weight+1
             focal[:,0] += h//8
             focal[:,1] += w//8
             print(focal*8)
@@ -105,6 +111,8 @@ def sample_dynamic(model, x, sigmas, extra_args=None, callback=None, disable=Non
         dt = sigmas[i + 1] - sigma_hat
         # Euler method
         x = x + d * dt
+        if early_terminate:
+            break
     summed = torch.nn.functional.conv2d(hist, kernel, groups=1)
     dump_image(hist)
     return x