fix next timestep offset in loop

2024-01-10 22:06:59 -06:00
parent ad575dcf87
commit f30f6cb91a
1 changed files with 3 additions and 3 deletions
@@ -268,9 +268,9 @@ class StyleAlignedReferenceSampler:
        # Replace first latent with the corresponding reference latent after each step
        def callback(step: int, x0: T, x: T, steps: int):
            preview_callback(step, x0, x, steps)
-            if (step < steps):
-                x[0] = ref_latents[step]
-                x0[0] = ref_latents[step]
+            if (step + 1 < steps):
+                x[0] = ref_latents[step+1]
+                x0[0] = ref_latents[step+1]

        # Register shared norms
        share_group_norm = share_norm in ["group", "both"]