From 450b65818476a6daae593be934028ff272d37121 Mon Sep 17 00:00:00 2001
From: Jedrzej Kosinski <kosinkadink1@gmail.com>
Date: Sun, 5 Jan 2025 15:30:54 -0600
Subject: [PATCH] Fixed backwards compatibility for scheduling, changed
 "sigmas" to "sample_sigmas" to match up with incoming ComfyUI PR

---
 animatediff/context.py         | 4 ++--
 animatediff/context_extras.py  | 4 ++--
 animatediff/model_injection.py | 4 ++--
 animatediff/sample_settings.py | 4 ++--
 animatediff/utils_model.py     | 1 +
 5 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/animatediff/context.py b/animatediff/context.py
index 106937d..df139e2 100644
--- a/animatediff/context.py
+++ b/animatediff/context.py
@@ -12,7 +12,7 @@ from comfy.model_base import BaseModel
 from comfy.model_patcher import ModelPatcher
 
 from .context_extras import ContextExtrasGroup
-from .utils_model import BIGMAX
+from .utils_model import BIGMAX_TENSOR
 from .utils_motion import get_sorted_list_via_attr
 
 
@@ -158,7 +158,7 @@ class ContextOptionsGroup:
         if curr_t == self._previous_t:
             return
         prev_index = self._current_index
-        max_sigma = torch.max(transformer_options.get("sigmas", BIGMAX))
+        max_sigma = torch.max(transformer_options.get("sample_sigmas", BIGMAX_TENSOR))
         # if met guaranteed steps, look for next context in case need to switch
         if self._current_used_steps >= self._current_context.get_effective_guarantee_steps(max_sigma):
             # if has next index, loop through and see if need to switch
diff --git a/animatediff/context_extras.py b/animatediff/context_extras.py
index dc7d959..de8be7b 100644
--- a/animatediff/context_extras.py
+++ b/animatediff/context_extras.py
@@ -5,7 +5,7 @@ from torch import Tensor
 
 from comfy.model_base import BaseModel
 
-from .utils_model import BIGMAX
+from .utils_model import BIGMAX_TENSOR
 from .utils_motion import (prepare_mask_batch, extend_to_batch_size, get_combined_multival, resize_multival,
                            get_sorted_list_via_attr)
 
@@ -345,7 +345,7 @@ class NaiveReuseKeyframeGroup:
         if curr_t == self._previous_t:
             return
         prev_index = self._current_index
-        max_sigma = torch.max(transformer_options.get("sigmas", BIGMAX))
+        max_sigma = torch.max(transformer_options.get("sample_sigmas", BIGMAX_TENSOR))
         # if met guaranteed steps, look for next keyframe in case need to switch
         if self._current_used_steps >= self._current_keyframe.get_effective_guarantee_steps(max_sigma):
             # if has next index, loop through and see if need t oswitch
diff --git a/animatediff/model_injection.py b/animatediff/model_injection.py
index 6ce1144..84f582c 100644
--- a/animatediff/model_injection.py
+++ b/animatediff/model_injection.py
@@ -30,7 +30,7 @@ from .utils_motion import (ADKeyframe, ADKeyframeGroup, MotionCompatibilityError
                            get_combined_multival, get_combined_input, get_combined_input_effect_multival,
                            ade_broadcast_image_to, extend_to_batch_size, prepare_mask_batch)
 from .motion_lora import MotionLoraInfo, MotionLoraList
-from .utils_model import get_motion_lora_path, get_motion_model_path, get_sd_model_type, vae_encode_raw_batched, BIGMAX
+from .utils_model import get_motion_lora_path, get_motion_model_path, get_sd_model_type, vae_encode_raw_batched, BIGMAX_TENSOR
 from .sample_settings import SampleSettings, SeedNoiseGeneration
 from .dinklink import DinkLinkConst, get_dinklink, get_acn_outer_sample_wrapper
 
@@ -333,7 +333,7 @@ class MotionModelAttachment:
         if curr_t == self.previous_t:
             return
         prev_index = self.current_index
-        max_sigma = torch.max(transformer_options.get("sigmas", BIGMAX))
+        max_sigma = torch.max(transformer_options.get("sample_sigmas", BIGMAX_TENSOR))
         # if met guaranteed steps, look for next keyframe in case need to switch
         if self.current_keyframe is None or self.current_used_steps >= self.current_keyframe.get_effective_guarantee_steps(max_sigma):
             # if has next index, loop through and see if need to switch
diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py
index d504875..5f0adde 100644
--- a/animatediff/sample_settings.py
+++ b/animatediff/sample_settings.py
@@ -14,7 +14,7 @@ from comfy.sd import VAE
 
 from . import freeinit
 from .context import ContextOptions, ContextOptionsGroup
-from .utils_model import SigmaSchedule, BIGMAX
+from .utils_model import SigmaSchedule, BIGMAX_TENSOR
 from .utils_motion import extend_to_batch_size, get_sorted_list_via_attr, prepare_mask_batch
 from .logger import logger
 
@@ -672,7 +672,7 @@ class CustomCFGKeyframeGroup:
         if curr_t == self._previous_t:
             return
         prev_index = self._current_index
-        max_sigma = torch.max(transformer_options.get("sigmas", BIGMAX))
+        max_sigma = torch.max(transformer_options.get("sample_sigmas", BIGMAX_TENSOR))
         # if met guaranteed steps, look for next keyframe in case need to switch
         if self._current_used_steps >= self._current_keyframe.get_effective_guarantee_steps(max_sigma):
             # if has next index, loop through and see if need t oswitch
diff --git a/animatediff/utils_model.py b/animatediff/utils_model.py
index 54bef00..d6659bc 100644
--- a/animatediff/utils_model.py
+++ b/animatediff/utils_model.py
@@ -23,6 +23,7 @@ from .logger import logger
 
 BIGMIN = -(2**53-1)
 BIGMAX = (2**53-1)
+BIGMAX_TENSOR = torch.tensor(BIGMAX)
 
 MAX_RESOLUTION = 16384  # mirrors ComfyUI's nodes.py MAX_RESOLUTION