From c4aa684cb3f4dc421ec8ac1ff07619b65cd6ae4b Mon Sep 17 00:00:00 2001
From: Jedrzej Kosinski <kosinkadink1@gmail.com>
Date: Sun, 4 Aug 2024 16:23:45 -0500
Subject: [PATCH] Fixed TimestepKeyframes being evaluated each time get_control
 is run instead of just once per step (technically could change results of
 some workflows, but now works as intended)

---
 adv_control/utils.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/adv_control/utils.py b/adv_control/utils.py
index 02ed6e7..30746fe 100644
--- a/adv_control/utils.py
+++ b/adv_control/utils.py
@@ -573,7 +573,8 @@ class AdvancedControlBase:
         self.full_latent_length = 0
         self.context_length = 0
         # timesteps
-        self.t: Tensor = None
+        self.t: float = None
+        self.prev_t: float = None
         self.batched_number: Union[int, IntWithCondOrUncond] = None
         self.batch_size: int = 0
         # weights + override
@@ -635,6 +636,9 @@ class AdvancedControlBase:
         self.t = float(t[0])
         self.batched_number = batched_number
         self.batch_size = len(t)
+        # check if t has changed (otherwise do nothing, as step already accounted for)
+        if self.t == self.prev_t:
+            return
         # get current step percent
         curr_t: float = self.t
         prev_index = self._current_timestep_index
@@ -670,7 +674,8 @@ class AdvancedControlBase:
                     # if eval_tk is outside of percent range, stop looking further
                     else:
                         break
-        
+        # update prev_t
+        self.prev_t = self.t
         # update steps current keyframe is used
         self._current_used_steps += 1
         # if index changed, apply overrides
@@ -936,6 +941,7 @@ class AdvancedControlBase:
         self.full_latent_length = 0
         self.context_length = 0
         self.t = None
+        self.prev_t = None
         self.batched_number = None
         self.batch_size = 0
         self.weights = None