From fd5d0ef6a103d35b951ebf6df94529e40203c31b Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Thu, 15 Aug 2024 05:01:13 -0600
Subject: [PATCH] Fix noise cache update to work better when the latent size
 changes

---
 py/sampling.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/py/sampling.py b/py/sampling.py
index d21f2ba..305b2b7 100644
--- a/py/sampling.py
+++ b/py/sampling.py
@@ -91,7 +91,6 @@ def composable_sampler(
             ss.update(0, step=step, substep=0)
             if step != 0:
                 nsc.reset_cache()
-                nsc.update_x(x)
                 ss.hist.reset()
                 for ms in merge_samplers:
                     ms.reset()
@@ -106,6 +105,8 @@ def composable_sampler(
             for idx in range(len(chunk_sigmas) - 1):
                 if idx > 0:
                     ss.update(idx, step=step, substep=0)
+                if step != 0:
+                    nsc.update_x(x)
                 # print(
                 #     f"STEP {step + 1:>3}: {ss.sigma.item():.03} -> {ss.sigma_next.item():.03} || up={ss.sigma_up.item():.03}, down={ss.sigma_down.item():.03}"
                 # )