more extensions

2026-10-15 00:03:06 +00:00 · 2024-02-20 08:30:59 -08:00
parent 8048a05d12
commit 7053eec79c
145 changed files with 7747 additions and 2775 deletions
@@ -46,7 +46,8 @@ NOTE: you can also use custom locations for models/motion loras by making use of
 - Infinite animation length support via sliding context windows across whole unet (Context Options) and/or within motion module (View Options)
 - Scheduling Context Options to change across different points in the sampling process
 - FreeInit and FreeNoise support (FreeInit is under iteration opts, FreeNoise is in SampleSettings' noise_type dropdown)
- Mixable Motion LoRAs from [original AnimateDiff repository](https://github.com/guoyww/animatediff/) implemented. Caveat: only really work on v2-based motion models like ```mm_sd_v15_v2```, ```mm-p_0.5.pth```, and ```mm-p_0.75.pth```
+- Mixable Motion LoRAs from [original AnimateDiff repository](https://github.com/guoyww/animatediff/) implemented. Caveat: the original loras really only work on v2-based motion models like ```mm_sd_v15_v2```, ```mm-p_0.5.pth```, and ```mm-p_0.75.pth```.
+     - UPDATE: New motion LoRAs without the v2 limitation can now be trained via the [AnimateDiff-MotionDirector repo](https://github.com/ExponentialML/AnimateDiff-MotionDirector). Shoutout to ExponentialML for implementing MotionDirector for AnimateDiff purposes!
 - Prompt travel using BatchPromptSchedule node from [ComfyUI_FizzNodes](https://github.com/FizzleDorf/ComfyUI_FizzNodes)
 - Scale and Effect multival inputs to control motion amount and motion model influence on generation.
     - Can be float, list of floats, or masks
@@ -57,15 +58,17 @@ NOTE: you can also use custom locations for models/motion loras by making use of
     - NOTE: You will need to use ```autoselect``` or ```linear (HotshotXL/default)``` beta_schedule, the sweetspot for context_length or total frames (when not using context) is 8 frames, and you will need to use an SDXL checkpoint.
 - AnimateDiff-SDXL support, with corresponding model. Currently, a beta version is out, which you can find info about at [AnimateDiff](https://github.com/guoyww/AnimateDiff/).
     - NOTE: You will need to use ```autoselect``` or ```linear (AnimateDiff-SDXL)``` beta_schedule. Other than that, same rules of thumb apply to AnimateDiff-SDXL as AnimateDiff.
+- [AnimateLCM](https://github.com/G-U-N/AnimateLCM) support
+     - NOTE: You will need to use ```autoselect``` or ```lcm``` or ```lcm[100_ots]``` beta_schedule. To use fully with LCM, be sure to use appropriate LCM lora, use the ```lcm``` sampler_name in KSampler nodes, and lower cfg to somewhere around 1.0 to 2.0. Don't forget to decrease steps (minimum = ~4 steps), since LCM converges faster (less steps). Increase step count to increase detail as desired.
 - AnimateDiff Keyframes to change Scale and Effect at different points in the sampling process.
 - fp8 support; requires newest ComfyUI and torch >= 2.1 (decreases VRAM usage, but changes outputs)
 - Mac M1/M2/M3 support
 - Usage of Context Options and Sample Settings outside of AnimateDiff via Gen2 Use Evolved Sampling node

 ## Upcoming Features
+- Maskable Motion LoRA
 - Maskable SD LoRA (and perhaps maskable SD Models as well)
 - [PIA](https://github.com/open-mmlab/PIA) support
- Motion LoRA training (experimental)
 - Anything else AnimateDiff-related that comes out


@@ -10,8 +10,10 @@ from .utils_motion import get_sorted_list_via_attr
 class ContextFuseMethod:
    FLAT = "flat"
    PYRAMID = "pyramid"
+    RELATIVE = "relative"

    LIST = [PYRAMID, FLAT]
+    LIST_STATIC = [PYRAMID, RELATIVE, FLAT]


 class ContextType:
@@ -331,6 +333,7 @@ def create_weights_pyramid(length: int, **kwargs) -> list[float]:
 FUSE_MAPPING = {
    ContextFuseMethod.FLAT: create_weights_flat,
    ContextFuseMethod.PYRAMID: create_weights_pyramid,
+    ContextFuseMethod.RELATIVE: create_weights_pyramid,
 }


@@ -13,7 +13,7 @@ from comfy.model_base import BaseModel

 from .ad_settings import AnimateDiffSettings
 from .context import ContextOptions, ContextOptions, ContextOptionsGroup
-from .motion_module_ad import AnimateDiffModel, has_mid_block, normalize_ad_state_dict
+from .motion_module_ad import AnimateDiffModel, AnimateDiffFormat, has_mid_block, normalize_ad_state_dict
 from .logger import logger
 from .utils_motion import ADKeyframe, ADKeyframeGroup, MotionCompatibilityError, get_combined_multival, normalize_min_max
 from .motion_lora import MotionLoraInfo, MotionLoraList
@@ -365,7 +365,8 @@ def load_motion_module_gen1(model_name: str, model: ModelPatcher, motion_lora: M
    ad_wrapper = AnimateDiffModel(mm_state_dict=mm_state_dict, mm_info=mm_info)
    ad_wrapper.to(model.model_dtype())
    ad_wrapper.to(model.offload_device)
-    load_result = ad_wrapper.load_state_dict(mm_state_dict)
+    is_animatelcm = mm_info.mm_format==AnimateDiffFormat.ANIMATELCM
+    load_result = ad_wrapper.load_state_dict(mm_state_dict, strict=not is_animatelcm)
    # TODO: report load_result of motion_module loading?
    # wrap motion_module into a ModelPatcher, to allow motion lora patches
    motion_model = MotionModelPatcher(model=ad_wrapper, load_device=model.load_device, offload_device=model.offload_device)
@@ -389,7 +390,11 @@ def load_motion_module_gen2(model_name: str, motion_model_settings: AnimateDiffS
    ad_wrapper = AnimateDiffModel(mm_state_dict=mm_state_dict, mm_info=mm_info)
    ad_wrapper.to(comfy.model_management.unet_dtype())
    ad_wrapper.to(comfy.model_management.unet_offload_device())
-    load_result = ad_wrapper.load_state_dict(mm_state_dict)
+    is_animatelcm = mm_info.mm_format==AnimateDiffFormat.ANIMATELCM
+    load_result = ad_wrapper.load_state_dict(mm_state_dict, strict=not is_animatelcm)
+    # TODO: manually check load_results for AnimateLCM models
+    if is_animatelcm:
+        pass
    # TODO: report load_result of motion_module loading?
    # wrap motion_module into a ModelPatcher, to allow motion lora patches
    motion_model = MotionModelPatcher(model=ad_wrapper, load_device=comfy.model_management.get_torch_device(),
@@ -15,7 +15,7 @@ from comfy.utils import repeat_to_batch_size
 import comfy.ops
 import comfy.model_management

-from .context import ContextOptions, get_context_weights, get_context_windows
+from .context import ContextFuseMethod, ContextOptions, get_context_weights, get_context_windows
 from .utils_motion import CrossAttentionMM, MotionCompatibilityError, extend_to_batch_size, prepare_mask_batch
 from .utils_model import BetaSchedules, ModelTypeSD
 from .logger import logger
@@ -31,6 +31,7 @@ def zero_module(module):
 class AnimateDiffFormat:
    ANIMATEDIFF = "AnimateDiff"
    HOTSHOTXL = "HotshotXL"
+    ANIMATELCM = "AnimateLCM"


 class AnimateDiffVersion:
@@ -58,6 +59,14 @@ def is_hotshotxl(mm_state_dict: dict[str, Tensor]) -> bool:
    return False


+def is_animatelcm(mm_state_dict: dict[str, Tensor]) -> bool:
+    # use lack of ANY pos_encoder keys to determine if animatelcm model
+    for key in mm_state_dict.keys():
+        if "pos_encoder" in key:
+            return False
+    return True
+
+
 def get_down_block_max(mm_state_dict: dict[str, Tensor]) -> int:
    # keep track of biggest down_block count in module
    biggest_block = 0
@@ -81,11 +90,14 @@ def has_mid_block(mm_state_dict: dict[str, Tensor]):
    return False


-def get_position_encoding_max_len(mm_state_dict: dict[str, Tensor], mm_name: str) -> int:
+def get_position_encoding_max_len(mm_state_dict: dict[str, Tensor], mm_name: str, mm_format: str) -> Union[int, None]:
    # use pos_encoder.pe entries to determine max length - [1, {max_length}, {320|640|1280}]
    for key in mm_state_dict.keys():
        if key.endswith("pos_encoder.pe"):
            return mm_state_dict[key].size(1) # get middle dim
+    # AnimateLCM models should have no pos_encoder entries, and assumed to be 64
+    if mm_format == AnimateDiffFormat.ANIMATELCM:
+        return 64
    raise MotionCompatibilityError(f"No pos_encoder.pe found in mm_state_dict - {mm_name} is not a valid AnimateDiff motion module!")


@@ -98,6 +110,11 @@ def find_hotshot_module_num(key: str) -> Union[int, None]:


 def normalize_ad_state_dict(mm_state_dict: dict[str, Tensor], mm_name: str) -> Tuple[dict[str, Tensor], AnimateDiffInfo]:
+    # from pathlib import Path
+    # with open(Path(__file__).parent.parent.parent / f"keys_{mm_name}.txt", "w") as afile:
+    #     for key, value in mm_state_dict.items():
+    #         afile.write(f"{key}:\t{value.shape}\n")
+    
    # remove all non-temporal keys (in case model has extra stuff in it)
    for key in list(mm_state_dict.keys()):
        if "temporal" not in key:
@@ -115,11 +132,13 @@ def normalize_ad_state_dict(mm_state_dict: dict[str, Tensor], mm_name: str) -> T
    mm_format = AnimateDiffFormat.ANIMATEDIFF
    if is_hotshotxl(mm_state_dict):
        mm_format = AnimateDiffFormat.HOTSHOTXL
+    if is_animatelcm(mm_state_dict):
+        mm_format = AnimateDiffFormat.ANIMATELCM
    # determine the model's version
    mm_version = AnimateDiffVersion.V1
    if has_mid_block(mm_state_dict):
        mm_version = AnimateDiffVersion.V2
-    elif sd_type==ModelTypeSD.SD1_5 and get_position_encoding_max_len(mm_state_dict, mm_name)==32:
+    elif sd_type==ModelTypeSD.SD1_5 and get_position_encoding_max_len(mm_state_dict, mm_name, mm_format)==32:
        mm_version = AnimateDiffVersion.V3
    info = AnimateDiffInfo(sd_type=sd_type, mm_format=mm_format, mm_version=mm_version, mm_name=mm_name)
    # convert to AnimateDiff format, if needed
@@ -157,7 +176,8 @@ class AnimateDiffModel(nn.Module):
        self.down_blocks: Iterable[MotionModule] = nn.ModuleList([])
        self.up_blocks: Iterable[MotionModule] = nn.ModuleList([])
        self.mid_block: Union[MotionModule, None] = None
-        self.encoding_max_len = get_position_encoding_max_len(mm_state_dict, mm_info.mm_name)
+        self.encoding_max_len = get_position_encoding_max_len(mm_state_dict, mm_info.mm_name, mm_info.mm_format)
+        self.has_position_encoding = self.encoding_max_len is not None
        # determine ops to use (to support fp8 properly)
        if comfy.model_management.unet_manual_cast(comfy.model_management.unet_dtype(), comfy.model_management.get_torch_device()) is None:
            ops = comfy.ops.disable_weight_init
@@ -170,20 +190,31 @@ class AnimateDiffModel(nn.Module):
            layer_channels = (320, 640, 1280, 1280)
        # fill out down/up blocks and middle block, if present
        for c in layer_channels:
-            self.down_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.DOWN, ops=ops))
+            self.down_blocks.append(MotionModule(c, temporal_position_encoding=self.has_position_encoding,
+                                                 temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.DOWN, ops=ops))
        for c in reversed(layer_channels):
-            self.up_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.UP, ops=ops))
+            self.up_blocks.append(MotionModule(c, temporal_position_encoding=self.has_position_encoding,
+                                               temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.UP, ops=ops))
        if has_mid_block(mm_state_dict):
-            self.mid_block = MotionModule(1280, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.MID, ops=ops)
+            self.mid_block = MotionModule(1280, temporal_position_encoding=self.has_position_encoding,
+                                          temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.MID, ops=ops)
        self.AD_video_length: int = 24

    def get_device_debug(self):
        return self.down_blocks[0].motion_modules[0].temporal_transformer.proj_in.weight.device

+    def is_length_valid_for_encoding_max_len(self, length: int):
+        if self.encoding_max_len is None:
+            return True
+        return length <= self.encoding_max_len
+
    def get_best_beta_schedule(self, log=False) -> str:
        to_return = None
        if self.mm_info.sd_type == ModelTypeSD.SD1_5:
-            to_return = BetaSchedules.SQRT_LINEAR
+            if self.mm_info.mm_format == AnimateDiffFormat.ANIMATELCM:
+                to_return = BetaSchedules.LCM  # while LCM_100 is the intended schedule, I find LCM to have much less flicker
+            else:
+                to_return = BetaSchedules.SQRT_LINEAR
        elif self.mm_info.sd_type == ModelTypeSD.SDXL:
            if self.mm_info.mm_format == AnimateDiffFormat.HOTSHOTXL:
                to_return = BetaSchedules.LINEAR
@@ -352,22 +383,28 @@ class AnimateDiffModel(nn.Module):


 class MotionModule(nn.Module):
-    def __init__(self, in_channels, temporal_position_encoding_max_len=24, block_type: str=BlockType.DOWN, ops=comfy.ops.disable_weight_init):
+    def __init__(self,
+            in_channels,
+            temporal_position_encoding=True,
+            temporal_position_encoding_max_len=24,
+            block_type: str=BlockType.DOWN,
+            ops=comfy.ops.disable_weight_init
+        ):
        super().__init__()
        if block_type == BlockType.MID:
            # mid blocks contain only a single VanillaTemporalModule
-            self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList([get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops)])
+            self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList([get_motion_module(in_channels, temporal_position_encoding, temporal_position_encoding_max_len, ops=ops)])
        else:
            # down blocks contain two VanillaTemporalModules
            self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList(
                [
-                    get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops),
-                    get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops)
+                    get_motion_module(in_channels, temporal_position_encoding, temporal_position_encoding_max_len, ops=ops),
+                    get_motion_module(in_channels, temporal_position_encoding, temporal_position_encoding_max_len, ops=ops)
                ]
            )
            # up blocks contain one additional VanillaTemporalModule
            if block_type == BlockType.UP: 
-                self.motion_modules.append(get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops))
+                self.motion_modules.append(get_motion_module(in_channels, temporal_position_encoding, temporal_position_encoding_max_len, ops=ops))
    
    def set_video_length(self, video_length: int, full_length: int):
        for motion_module in self.motion_modules:
@@ -398,8 +435,8 @@ class MotionModule(nn.Module):
            motion_module.reset_temp_vars()


-def get_motion_module(in_channels, temporal_position_encoding_max_len, ops=comfy.ops.disable_weight_init):
-    return VanillaTemporalModule(in_channels=in_channels, temporal_position_encoding_max_len=temporal_position_encoding_max_len, ops=ops)
+def get_motion_module(in_channels, temporal_position_encoding, temporal_position_encoding_max_len, ops=comfy.ops.disable_weight_init):
+    return VanillaTemporalModule(in_channels=in_channels, temporal_position_encoding=temporal_position_encoding, temporal_position_encoding_max_len=temporal_position_encoding_max_len, ops=ops)


 class VanillaTemporalModule(nn.Module):
@@ -776,11 +813,9 @@ class TemporalTransformerBlock(nn.Module):
            hidden_states = rearrange(hidden_states, "(b f) d c -> b f d c", f=video_length)
            value_final = torch.zeros_like(hidden_states)
            count_final = torch.zeros_like(hidden_states)
+            # bias_final = [0.0] * video_length
            batched_conds = hidden_states.size(1) // video_length
            for sub_idxs in views:
-                weights = get_context_weights(len(sub_idxs), view_options.fuse_method) * batched_conds
-                weights_tensor = torch.Tensor(weights).to(device=hidden_states.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
-
                sub_hidden_states = rearrange(hidden_states[:, sub_idxs], "b f d c -> (b f) d c")
                for attention_block, norm in zip(self.attention_blocks, self.norms):
                    norm_hidden_states = norm(sub_hidden_states).to(sub_hidden_states.dtype)
@@ -797,14 +832,35 @@ class TemporalTransformerBlock(nn.Module):
                    )
                sub_hidden_states = rearrange(sub_hidden_states, "(b f) d c -> b f d c", f=len(sub_idxs))

+                # if view_options.fuse_method == ContextFuseMethod.RELATIVE:
+                #     for pos, idx in enumerate(sub_idxs):
+                #         # bias is the influence of a specific index in relation to the whole context window
+                #         bias = 1 - abs(idx - (sub_idxs[0] + sub_idxs[-1]) / 2) / ((sub_idxs[-1] - sub_idxs[0] + 1e-2) / 2)
+                #         bias = max(1e-2, bias)
+                #         # take weighted averate relative to total bias of current idx
+                #         bias_total = bias_final[idx]
+                #         prev_weight = torch.tensor([bias_total / (bias_total + bias)],
+                #                                    dtype=value_final.dtype, device=value_final.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+                #         #prev_weight = torch.cat([prev_weight]*value_final.shape[1], dim=1)
+                #         new_weight = torch.tensor([bias / (bias_total + bias)],
+                #                                    dtype=value_final.dtype, device=value_final.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+                #         #new_weight = torch.cat([new_weight]*value_final.shape[1], dim=1)
+                #         test = value_final[:, idx:idx+1, :, :]
+                #         value_final[:, idx:idx+1, :, :] = value_final[:, idx:idx+1, :, :] * prev_weight + sub_hidden_states[:, pos:pos+1, : ,:] * new_weight
+                #         bias_final[idx] = bias_total + bias
+                # else:
+                weights = get_context_weights(len(sub_idxs), view_options.fuse_method) * batched_conds
+                weights_tensor = torch.Tensor(weights).to(device=hidden_states.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
                value_final[:, sub_idxs] += sub_hidden_states * weights_tensor
                count_final[:, sub_idxs] += weights_tensor
            
-            # get weighted average of sub_hidden_states
+            # get weighted average of sub_hidden_states, if fuse method requires it
+            # if view_options.fuse_method != ContextFuseMethod.RELATIVE:
            hidden_states = value_final / count_final
            hidden_states = rearrange(hidden_states, "b f d c -> (b f) d c")
            del value_final
            del count_final
+            # del bias_final

        hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states

@@ -1,65 +1,28 @@
-from pathlib import Path
-import torch
-
 import comfy.sample as comfy_sample
-from comfy.model_patcher import ModelPatcher

-from .logger import logger
-from .utils_model import BetaSchedules, get_available_motion_loras, get_available_motion_models, get_motion_lora_path
-from .motion_lora import MotionLoraInfo, MotionLoraList
-from .model_injection import InjectionParams, ModelPatcherAndInjector, AnimateDiffSettings, load_motion_module_gen1
-from .sample_settings import SampleSettings, SeedNoiseGeneration
 from .sampling import motion_sample_factory

 from .nodes_gen1 import (AnimateDiffLoaderGen1, LegacyAnimateDiffLoaderWithContext, AnimateDiffModelSettings,
                         AnimateDiffModelSettingsSimple, AnimateDiffModelSettingsAdvanced, AnimateDiffModelSettingsAdvancedAttnStrengths)
 from .nodes_gen2 import UseEvolvedSamplingNode, ApplyAnimateDiffModelNode, ApplyAnimateDiffModelBasicNode, LoadAnimateDiffModelNode, ADKeyframeNode
-from .nodes_multival import MultivalDynamicNode, MultivalFloatNode, MultivalScaledMaskNode
-from .nodes_sample import FreeInitOptionsNode, NoiseLayerAddWeightedNode, SampleSettingsNode, NoiseLayerAddNode, NoiseLayerReplaceNode, IterationOptionsNode
+from .nodes_multival import MultivalDynamicNode, MultivalScaledMaskNode
+from .nodes_sample import (FreeInitOptionsNode, NoiseLayerAddWeightedNode, SampleSettingsNode, NoiseLayerAddNode, NoiseLayerReplaceNode, IterationOptionsNode,
+                           CustomCFGNode, CustomCFGKeyframeNode)
+from .nodes_sigma_schedule import (SigmaScheduleNode, RawSigmaScheduleNode, WeightedAverageSigmaScheduleNode, InterpolatedWeightedAverageSigmaScheduleNode, SplitAndCombineSigmaScheduleNode)
 from .nodes_context import (LegacyLoopedUniformContextOptionsNode, LoopedUniformContextOptionsNode, LoopedUniformViewOptionsNode, StandardUniformContextOptionsNode, StandardStaticContextOptionsNode, BatchedContextOptionsNode,
                            StandardStaticViewOptionsNode, StandardUniformViewOptionsNode, ViewAsContextOptionsNode)
 from .nodes_ad_settings import AnimateDiffSettingsNode, ManualAdjustPENode, SweetspotStretchPENode, FullStretchPENode
 from .nodes_extras import AnimateDiffUnload, EmptyLatentImageLarge, CheckpointLoaderSimpleWithNoiseSelect
 from .nodes_deprecated import AnimateDiffLoader_Deprecated, AnimateDiffLoaderAdvanced_Deprecated, AnimateDiffCombine_Deprecated
+from .nodes_lora import AnimateDiffLoraLoader, MaskedLoraLoader
+
+from .logger import logger

 # override comfy_sample.sample with animatediff-support version
 comfy_sample.sample = motion_sample_factory(comfy_sample.sample)
 comfy_sample.sample_custom = motion_sample_factory(comfy_sample.sample_custom, is_custom=True)


-class AnimateDiffLoraLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "lora_name": (get_available_motion_loras(),),
-                "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}),
-            },
-            "optional": {
-                "prev_motion_lora": ("MOTION_LORA",),
-            }
-        }
-    
-    RETURN_TYPES = ("MOTION_LORA",)
-    CATEGORY = "Animate Diff 🎭🅐🅓"
-    FUNCTION = "load_motion_lora"
-
-    def load_motion_lora(self, lora_name: str, strength: float, prev_motion_lora: MotionLoraList=None):
-        if prev_motion_lora is None:
-            prev_motion_lora = MotionLoraList()
-        else:
-            prev_motion_lora = prev_motion_lora.clone()
-        # check if motion lora with name exists
-        lora_path = get_motion_lora_path(lora_name)
-        if not Path(lora_path).is_file():
-            raise FileNotFoundError(f"Motion lora with name '{lora_name}' not found.")
-        # create motion lora info to be loaded in AnimateDiff Loader
-        lora_info = MotionLoraInfo(name=lora_name, strength=strength)
-        prev_motion_lora.add_lora(lora_info)
-
-        return (prev_motion_lora,)
-
-
 NODE_CLASS_MAPPINGS = {
    # Unencapsulated
    "ADE_AnimateDiffLoRALoader": AnimateDiffLoraLoader,
@@ -91,6 +54,14 @@ NODE_CLASS_MAPPINGS = {
    "ADE_AdjustPESweetspotStretch": SweetspotStretchPENode,
    "ADE_AdjustPEFullStretch": FullStretchPENode,
    "ADE_AdjustPEManual": ManualAdjustPENode,
+    # Sample Settings
+    "ADE_CustomCFG": CustomCFGNode,
+    "ADE_CustomCFGKeyframe": CustomCFGKeyframeNode,
+    "ADE_SigmaSchedule": SigmaScheduleNode,
+    "ADE_RawSigmaSchedule": RawSigmaScheduleNode,
+    "ADE_SigmaScheduleWeightedAverage": WeightedAverageSigmaScheduleNode,
+    "ADE_SigmaScheduleWeightedAverageInterp": InterpolatedWeightedAverageSigmaScheduleNode,
+    "ADE_SigmaScheduleSplitAndCombine": SplitAndCombineSigmaScheduleNode,
    # Extras Nodes
    "ADE_AnimateDiffUnload": AnimateDiffUnload,
    "ADE_EmptyLatentImageLarge": EmptyLatentImageLarge,
@@ -107,6 +78,8 @@ NODE_CLASS_MAPPINGS = {
    "ADE_ApplyAnimateDiffModelSimple": ApplyAnimateDiffModelBasicNode,
    "ADE_ApplyAnimateDiffModel": ApplyAnimateDiffModelNode,
    "ADE_LoadAnimateDiffModel": LoadAnimateDiffModelNode,
+    # MaskedLoraLoader
+    #"ADE_MaskedLoadLora": MaskedLoraLoader,
    # Deprecated Nodes
    "AnimateDiffLoaderV1": AnimateDiffLoader_Deprecated,
    "ADE_AnimateDiffLoaderV1Advanced": AnimateDiffLoaderAdvanced_Deprecated,
@@ -143,6 +116,14 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "ADE_AdjustPESweetspotStretch": "Adjust PE [Sweetspot Stretch] 🎭🅐🅓",
    "ADE_AdjustPEFullStretch": "Adjust PE [Full Stretch] 🎭🅐🅓",
    "ADE_AdjustPEManual": "Adjust PE [Manual] 🎭🅐🅓",
+    # Sample Settings
+    "ADE_CustomCFG": "Custom CFG 🎭🅐🅓",
+    "ADE_CustomCFGKeyframe": "Custom CFG Keyframe 🎭🅐🅓",
+    "ADE_SigmaSchedule": "Create Sigma Schedule 🎭🅐🅓",
+    "ADE_RawSigmaSchedule": "Create Raw Sigma Schedule 🎭🅐🅓",
+    "ADE_SigmaScheduleWeightedAverage": "Sigma Schedule Weighted Mean 🎭🅐🅓",
+    "ADE_SigmaScheduleWeightedAverageInterp": "Sigma Schedule Interpolated Mean 🎭🅐🅓",
+    "ADE_SigmaScheduleSplitAndCombine": "Sigma Schedule Split Combine 🎭🅐🅓",
    # Extras Nodes
    "ADE_AnimateDiffUnload": "AnimateDiff Unload 🎭🅐🅓",
    "ADE_EmptyLatentImageLarge": "Empty Latent Image (Big Batch) 🎭🅐🅓",
@@ -159,8 +140,10 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "ADE_ApplyAnimateDiffModelSimple": "Apply AnimateDiff Model 🎭🅐🅓②",
    "ADE_ApplyAnimateDiffModel": "Apply AnimateDiff Model (Adv.) 🎭🅐🅓②",
    "ADE_LoadAnimateDiffModel": "Load AnimateDiff Model 🎭🅐🅓②",
+    # MaskedLoraLoader
+    #"ADE_MaskedLoadLora": "Load LoRA (Masked) 🎭🅐🅓",
    # Deprecated Nodes
    "AnimateDiffLoaderV1": "AnimateDiff Loader [DEPRECATED] 🎭🅐🅓",
    "ADE_AnimateDiffLoaderV1Advanced": "AnimateDiff Loader (Advanced) [DEPRECATED] 🎭🅐🅓",
-    "ADE_AnimateDiffCombine": "DO NOT USE, USE VideoCombine from ComfyUI-VideoHelperSuite instead! AnimateDiff Combine [DEPRECATED, DO NOT USE] 🎭🅐🅓",
+    "ADE_AnimateDiffCombine": "AnimateDiff Combine [DEPRECATED, Use Video Combine (VHS) Instead!] 🎭🅐🅓",
 }
@@ -145,7 +145,7 @@ class StandardStaticContextOptionsNode:
                "context_overlap": ("INT", {"default": 4, "min": 0, "max": OVERLAP_MAX}),
            },
            "optional": {
-                "fuse_method": (ContextFuseMethod.LIST,),
+                "fuse_method": (ContextFuseMethod.LIST_STATIC,),
                "use_on_equal_length": ("BOOLEAN", {"default": False},),
                "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "guarantee_steps": ("INT", {"default": 1, "min": 0, "max": BIGMAX}),
@@ -12,7 +12,7 @@ from PIL.PngImagePlugin import PngInfo
 import folder_paths
 from comfy.model_patcher import ModelPatcher

-from .context import ContextSchedules, ContextOptions
+from .context import ContextOptionsGroup, ContextOptions, ContextSchedules
 from .logger import logger
 from .utils_model import Folders, BetaSchedules, get_available_motion_models
 from .model_injection import ModelPatcherAndInjector, InjectionParams, MotionModelGroup, load_motion_module_gen1
@@ -109,16 +109,18 @@ class AnimateDiffLoaderAdvanced_Deprecated:
                model_name=model_name,
                apply_v2_properly=False,
        )
-        # set context settings
-        params.set_context(
+        context_group = ContextOptionsGroup()
+        context_group.add(
            ContextOptions(
                context_length=context_length,
                context_stride=context_stride,
                context_overlap=context_overlap,
                context_schedule=context_schedule,
                closed_loop=closed_loop,
+                )
            )
-        )
+        # set context settings
+        params.set_context(context_options=context_group)
        # inject for use in sampling code
        model = ModelPatcherAndInjector(model)
        model.motion_models = MotionModelGroup(motion_model)
@@ -6,13 +6,13 @@ from comfy.model_patcher import ModelPatcher
 from comfy.sd import load_checkpoint_guess_config

 from .logger import logger
-from .utils_model import IsChangedHelper, BetaSchedules
+from .utils_model import BetaSchedules
 from .model_injection import get_vanilla_model_patcher


 class AnimateDiffUnload:
    def __init__(self) -> None:
-        self.change = IsChangedHelper()
+        pass

    @classmethod
    def INPUT_TYPES(s):
@@ -81,13 +81,20 @@ class AnimateDiffLoaderGen1:
        model.sample_settings = sample_settings if sample_settings is not None else SampleSettings()
        model.motion_injection_params = params
        
-        # save model sampling from BetaSchedule as object patch
-        # if autoselect, get suggested beta_schedule from motion model
-        if beta_schedule == BetaSchedules.AUTOSELECT and not model.motion_models.is_empty():
-            beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
-        new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
-        if new_model_sampling is not None:
-            model.add_object_patch("model_sampling", new_model_sampling)
+        if model.sample_settings.custom_cfg is not None:
+            logger.info("[Sample Settings] custom_cfg is set; will override any KSampler cfg values or patches.")
+
+        if model.sample_settings.sigma_schedule is not None:
+            logger.info("[Sample Settings] sigma_schedule is set; will override beta_schedule.")
+            model.add_object_patch("model_sampling", model.sample_settings.sigma_schedule.clone().model_sampling)
+        else:
+            # save model sampling from BetaSchedule as object patch
+            # if autoselect, get suggested beta_schedule from motion model
+            if beta_schedule == BetaSchedules.AUTOSELECT and not model.motion_models.is_empty():
+                beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
+            new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
+            if new_model_sampling is not None:
+                model.add_object_patch("model_sampling", new_model_sampling)

        del motion_model
        return (model,)
@@ -59,13 +59,20 @@ class UseEvolvedSamplingNode:
        model.sample_settings = sample_settings if sample_settings is not None else SampleSettings()
        model.motion_injection_params = params

-        # save model_sampling from BetaSchedule as object patch
-        # if autoselect, get suggested beta_schedule from motion model
-        if beta_schedule == BetaSchedules.AUTOSELECT and not model.motion_models.is_empty():
-            beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
-        new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
-        if new_model_sampling is not None:
-            model.add_object_patch("model_sampling", new_model_sampling)
+        if model.sample_settings.custom_cfg is not None:
+            logger.info("[Sample Settings] custom_cfg is set; will override any KSampler cfg values or patches.")
+
+        if model.sample_settings.sigma_schedule is not None:
+            logger.info("[Sample Settings] sigma_schedule is set; will override beta_schedule.")
+            model.add_object_patch("model_sampling", model.sample_settings.sigma_schedule.clone().model_sampling)
+        else:
+            # save model_sampling from BetaSchedule as object patch
+            # if autoselect, get suggested beta_schedule from motion model
+            if beta_schedule == BetaSchedules.AUTOSELECT and not model.motion_models.is_empty():
+                beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
+            new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
+            if new_model_sampling is not None:
+                model.add_object_patch("model_sampling", new_model_sampling)
        
        del m_models
        return (model,)
@@ -4,7 +4,7 @@ from typing import Union
 import torch
 from torch import Tensor

-from .utils_motion import linear_conversion, normalize_min_max
+from .utils_motion import linear_conversion, normalize_min_max, extend_to_batch_size


 class ScaleType:
@@ -41,24 +41,64 @@ class MultivalDynamicNode:
                if len(float_val) < mask_optional.shape[0]:
                    # copies last entry enough times to match mask shape
                    float_val = float_val + float_val[-1]*(mask_optional.shape[0]-len(float_val))
+                if mask_optional.shape[0] < len(float_val):
+                    mask_optional = extend_to_batch_size(mask_optional, len(float_val))
                float_val = float_val[:mask_optional.shape[0]]
+            float_val: Tensor = torch.tensor(float_val).unsqueeze(-1).unsqueeze(-1)
        # now that inputs are normalized, figure out what value to actually return
        if mask_optional is not None:
            mask_optional = mask_optional.clone()
-            mask_optional = mask_optional * float_val
+            if float_is_iterable:
+                mask_optional = mask_optional[:] * float_val.to(mask_optional.dtype).to(mask_optional.device)
+            else:
+                mask_optional = mask_optional * float_val
            return (mask_optional,)
        else:
            if not float_is_iterable:
                return (float_val,)
            # create a dummy mask of b,h,w=float_len,1,1 (sigle pixel)
            # purpose is for float input to work with mask code, without special cases
-            float_len = len(float_val) if float_is_iterable else 1
+            float_len = float_val.shape[0] if float_is_iterable else 1
            shape = (float_len,1,1)
            mask_optional = torch.ones(shape)
-            mask_optional = mask_optional * float_val
+            mask_optional = mask_optional[:] * float_val.to(mask_optional.dtype).to(mask_optional.device)
            return (mask_optional,)


+class MultivalScaledMaskNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "min_float_val": ("FLOAT", {"default": 0.0, "min": 0.0, "step": 0.001}),
+                "max_float_val": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
+                "mask": ("MASK",),
+            },
+            "optional": {
+                "scaling": (ScaleType.LIST,),
+            }
+        }
+
+    RETURN_TYPES = ("MULTIVAL",)
+    CATEGORY = "Animate Diff 🎭🅐🅓/multival"
+    FUNCTION = "create_multival"
+
+    def create_multival(self, min_float_val: float, max_float_val: float, mask: Tensor, scaling: str=ScaleType.ABSOLUTE):
+        # TODO: allow min_float_val and max_float_val to be list[float]
+        if isinstance(min_float_val, Iterable):
+            raise ValueError(f"min_float_val must be type float (no lists allowed here), not {type(min_float_val).__name__}.")
+        if isinstance(max_float_val, Iterable):
+            raise ValueError(f"max_float_val must be type float (no lists allowed here), not {type(max_float_val).__name__}.")
+        
+        if scaling == ScaleType.ABSOLUTE:
+            mask = linear_conversion(mask.clone(), new_min=min_float_val, new_max=max_float_val)
+        elif scaling == ScaleType.RELATIVE:
+            mask = normalize_min_max(mask.clone(), new_min=min_float_val, new_max=max_float_val)
+        else:
+            raise ValueError(f"scaling '{scaling}' not recognized.")
+        return MultivalDynamicNode.create_multival(self, mask_optional=mask)
+
+
 class MultivalDynamicFloatInputNode:
    @classmethod
    def INPUT_TYPES(s):
@@ -94,36 +134,3 @@ class MultivalFloatNode:

    def create_multival(self, float_val: Union[float, list[float]]=None):
        return MultivalDynamicNode.create_multival(self, float_val=float_val)
-
-
-class MultivalScaledMaskNode:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "min_float_val": ("FLOAT", {"default": 0.0, "min": 0.0, "step": 0.001}),
-                "max_float_val": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
-                "mask": ("MASK",),
-            },
-            "optional": {
-                "scaling": (ScaleType.LIST,),
-            }
-        }
-
-    RETURN_TYPES = ("MULTIVAL",)
-    CATEGORY = "Animate Diff 🎭🅐🅓/multival"
-    FUNCTION = "create_multival"
-
-    def create_multival(self, min_float_val: float, max_float_val: float, mask: Tensor, scaling: str=ScaleType.ABSOLUTE):
-        if isinstance(min_float_val, Iterable):
-            raise ValueError(f"min_float_val must be type float (no lists allowed here), not {type(min_float_val).__name__}.")
-        if isinstance(max_float_val, Iterable):
-            raise ValueError(f"max_float_val must be type float (no lists allowed here), not {type(max_float_val).__name__}.")
-        
-        if scaling == ScaleType.ABSOLUTE:
-            mask = linear_conversion(mask.clone(), new_min=min_float_val, new_max=max_float_val)
-        elif scaling == ScaleType.RELATIVE:
-            mask = normalize_min_max(mask.clone(), new_min=min_float_val, new_max=max_float_val)
-        else:
-            raise ValueError(f"scaling '{scaling}' not recognized.")
-        return MultivalDynamicNode.create_multival(self, mask_optional=mask)
@@ -1,8 +1,11 @@
+from typing import Union
 from torch import Tensor

 from .freeinit import FreeInitFilter
-from .sample_settings import FreeInitOptions, IterationOptions, NoiseLayerAdd, NoiseLayerAddWeighted, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType, SeedNoiseGeneration, SampleSettings
-from .utils_model import BIGMIN, BIGMAX
+from .sample_settings import (FreeInitOptions, IterationOptions,
+                              NoiseLayerAdd, NoiseLayerAddWeighted, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType,
+                              SeedNoiseGeneration, SampleSettings, CustomCFGKeyframeGroup, CustomCFGKeyframe)
+from .utils_model import BIGMIN, BIGMAX, SigmaSchedule


 class SampleSettingsNode:
@@ -20,18 +23,22 @@ class SampleSettingsNode:
                "iteration_opts": ("ITERATION_OPTS",),
                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
                "adapt_denoise_steps": ("BOOLEAN", {"default": False},),
+                "custom_cfg": ("CUSTOM_CFG",),
+                "sigma_schedule": ("SIGMA_SCHEDULE",),
            }
        }
-    
+
    RETURN_TYPES = ("SAMPLE_SETTINGS",)
    RETURN_NAMES = ("settings",)
    CATEGORY = "Animate Diff 🎭🅐🅓"
    FUNCTION = "create_settings"

    def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, seed_offset: int, noise_layers: NoiseLayerGroup=None,
-                        iteration_opts: IterationOptions=None, seed_override: int=None, adapt_denoise_steps=False):
+                        iteration_opts: IterationOptions=None, seed_override: int=None, adapt_denoise_steps=False,
+                        custom_cfg: CustomCFGKeyframeGroup=None, sigma_schedule: SigmaSchedule=None):
        sampling_settings = SampleSettings(batch_offset=batch_offset, noise_type=noise_type, seed_gen=seed_gen, seed_offset=seed_offset, noise_layers=noise_layers,
-                                           iteration_opts=iteration_opts, seed_override=seed_override, adapt_denoise_steps=adapt_denoise_steps)
+                                           iteration_opts=iteration_opts, seed_override=seed_override, adapt_denoise_steps=adapt_denoise_steps,
+                                           custom_cfg=custom_cfg, sigma_schedule=sigma_schedule)
        return (sampling_settings,)


@@ -51,7 +58,7 @@ class NoiseLayerReplaceNode:
                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
            }
        }
-    
+
    RETURN_TYPES = ("NOISE_LAYERS",)
    CATEGORY = "Animate Diff 🎭🅐🅓/noise layers"
    FUNCTION = "create_layers"
@@ -86,7 +93,7 @@ class NoiseLayerAddNode:
                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
            }
        }
-    
+
    RETURN_TYPES = ("NOISE_LAYERS",)
    CATEGORY = "Animate Diff 🎭🅐🅓/noise layers"
    FUNCTION = "create_layers"
@@ -124,7 +131,7 @@ class NoiseLayerAddWeightedNode:
                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
            }
        }
-    
+
    RETURN_TYPES = ("NOISE_LAYERS",)
    CATEGORY = "Animate Diff 🎭🅐🅓/noise layers"
    FUNCTION = "create_layers"
@@ -156,7 +163,7 @@ class IterationOptionsNode:
                "iter_seed_offset": ("INT", {"default": 0, "min": BIGMIN, "max": BIGMAX}),
            }
        }
-    
+
    RETURN_TYPES = ("ITERATION_OPTS",)
    CATEGORY = "Animate Diff 🎭🅐🅓/iteration opts"
    FUNCTION = "create_iter_opts"
@@ -185,7 +192,7 @@ class FreeInitOptionsNode:
                "iter_seed_offset": ("INT", {"default": 1, "min": BIGMIN, "max": BIGMAX}),
            }
        }
-    
+
    RETURN_TYPES = ("ITERATION_OPTS",)
    CATEGORY = "Animate Diff 🎭🅐🅓/iteration opts"
    FUNCTION = "create_iter_opts"
@@ -198,3 +205,51 @@ class FreeInitOptionsNode:
                                    filter=filter, d_s=d_s, d_t=d_t, n=n_butterworth, init_type=init_type,
                                    iter_batch_offset=iter_batch_offset, iter_seed_offset=iter_seed_offset)
        return (iter_opts,)
+
+
+class CustomCFGNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "cfg_multival": ("MULTIVAL",),
+            }
+        }
+
+    RETURN_TYPES = ("CUSTOM_CFG",)
+    CATEGORY = "Animate Diff 🎭🅐🅓/sample settings"
+    FUNCTION = "create_custom_cfg"
+
+    def create_custom_cfg(self, cfg_multival: Union[float, Tensor]):
+        keyframe = CustomCFGKeyframe(cfg_multival=cfg_multival)
+        cfg_custom = CustomCFGKeyframeGroup()
+        cfg_custom.add(keyframe)
+        return (cfg_custom,)
+
+
+class CustomCFGKeyframeNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "cfg_multival": ("MULTIVAL",),
+                "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                "guarantee_steps": ("INT", {"default": 1, "min": 0, "max": BIGMAX}),
+            },
+            "optional": {
+                "prev_custom_cfg": ("CUSTOM_CFG",),
+            }
+        }
+
+    RETURN_TYPES = ("CUSTOM_CFG",)
+    CATEGORY = "Animate Diff 🎭🅐🅓/sample settings"
+    FUNCTION = "create_custom_cfg"
+
+    def create_custom_cfg(self, cfg_multival: Union[float, Tensor], start_percent: float=0.0, guarantee_steps: int=1,
+                          prev_custom_cfg: CustomCFGKeyframeGroup=None):
+        if not prev_custom_cfg:
+            prev_custom_cfg = CustomCFGKeyframeGroup()
+        prev_custom_cfg = prev_custom_cfg.clone()
+        keyframe = CustomCFGKeyframe(cfg_multival=cfg_multival, start_percent=start_percent, guarantee_steps=guarantee_steps)
+        prev_custom_cfg.add(keyframe)
+        return (prev_custom_cfg,)
@@ -1,13 +1,17 @@
 from collections.abc import Iterable
+from typing import Union
 import torch
 from torch import Tensor

 import comfy.sample
 import comfy.samplers
 from comfy.model_patcher import ModelPatcher
+from comfy.model_base import BaseModel

 from . import freeinit
 from .context import ContextOptions, ContextOptionsGroup
+from .utils_model import SigmaSchedule
+from .utils_motion import extend_to_batch_size, get_sorted_list_via_attr, prepare_mask_batch
 from .logger import logger


@@ -48,7 +52,8 @@ class NoiseNormalize:

 class SampleSettings:
    def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, seed_offset: int=0, noise_layers: 'NoiseLayerGroup'=None,
-                 iteration_opts=None, seed_override:int=None, negative_cond_flipflop=False, adapt_denoise_steps: bool=False):
+                 iteration_opts=None, seed_override:int=None, negative_cond_flipflop=False, adapt_denoise_steps: bool=False,
+                 custom_cfg: 'CustomCFGKeyframeGroup'=None, sigma_schedule: SigmaSchedule=None):
        self.batch_offset = batch_offset
        self.noise_type = noise_type if noise_type is not None else NoiseLayerType.DEFAULT
        self.seed_gen = seed_gen if seed_gen is not None else SeedNoiseGeneration.COMFY
@@ -58,6 +63,8 @@ class SampleSettings:
        self.seed_override = seed_override
        self.negative_cond_flipflop = negative_cond_flipflop
        self.adapt_denoise_steps = adapt_denoise_steps
+        self.custom_cfg = custom_cfg.clone() if custom_cfg else custom_cfg
+        self.sigma_schedule = sigma_schedule
    
    def prepare_noise(self, seed: int, latents: Tensor, noise: Tensor, extra_seed_offset=0, extra_args:dict={}, force_create_noise=True):
        if self.seed_override is not None:
@@ -82,10 +89,18 @@ class SampleSettings:
        # noise prepared now
        return noise
    
+    def pre_run(self, model: ModelPatcher):
+        if self.custom_cfg is not None:
+            self.custom_cfg.reset()
+    
+    def cleanup(self):
+        if self.custom_cfg is not None:
+            self.custom_cfg.reset()
+
    def clone(self):
        return SampleSettings(batch_offset=self.batch_offset, noise_type=self.noise_type, seed_gen=self.seed_gen, seed_offset=self.seed_offset,
                           noise_layers=self.noise_layers.clone(), iteration_opts=self.iteration_opts, seed_override=self.seed_override,
-                           negative_cond_flipflop=self.negative_cond_flipflop, adapt_denoise_steps=self.adapt_denoise_steps)
+                           negative_cond_flipflop=self.negative_cond_flipflop, adapt_denoise_steps=self.adapt_denoise_steps, custom_cfg=self.custom_cfg, sigma_schedule=self.sigma_schedule)


 class NoiseLayer:
@@ -435,3 +450,106 @@ class FreeInitOptions(IterationOptions):
            return cached_latents, noised_latents
        else:
            raise ValueError(f"FreeInit init_type '{self.init_type}' is not recognized.")
+
+
+class CustomCFGKeyframe:
+    def __init__(self, cfg_multival: Union[float, Tensor], start_percent=0.0, guarantee_steps=1):
+        self.cfg_multival = cfg_multival
+        # scheduling
+        self.start_percent = float(start_percent)
+        self.start_t = 999999999.9
+        self.guarantee_steps = guarantee_steps
+    
+    def clone(self):
+        c = CustomCFGKeyframe(cfg_multival=self.cfg_multival,
+                              start_percent=self.start_percent, guarantee_steps=self.guarantee_steps)
+        c.start_t = self.start_t
+        return c
+
+
+class CustomCFGKeyframeGroup:
+    def __init__(self):
+        self.keyframes: list[CustomCFGKeyframe] = []
+        self._current_keyframe: CustomCFGKeyframe = None
+        self._current_used_steps: int = 0
+        self._current_index: int = 0
+    
+    def reset(self):
+        self._current_keyframe = None
+        self._current_used_steps = 0
+        self._current_index = 0
+        self._set_first_as_current()
+
+    def add(self, keyframe: CustomCFGKeyframe):
+        # add to end of list, then sort
+        self.keyframes.append(keyframe)
+        self.keyframes = get_sorted_list_via_attr(self.keyframes, "start_percent")
+        self._set_first_as_current()
+
+    def _set_first_as_current(self):
+        if len(self.keyframes) > 0:
+            self._current_keyframe = self.keyframes[0]
+        else:
+            self._current_keyframe = None
+
+    def has_index(self, index: int) -> int:
+        return index >=0 and index < len(self.keyframes)
+
+    def is_empty(self) -> bool:
+        return len(self.keyframes) == 0
+    
+    def clone(self):
+        cloned = CustomCFGKeyframeGroup()
+        for keyframe in self.keyframes:
+            cloned.keyframes.append(keyframe)
+        cloned._set_first_as_current()
+        return cloned
+    
+    def initialize_timesteps(self, model: BaseModel):
+        for keyframe in self.keyframes:
+            keyframe.start_t = model.model_sampling.percent_to_sigma(keyframe.start_percent)
+    
+    def prepare_current_keyframe(self, t: Tensor):
+        curr_t: float = t[0]
+        prev_index = self._current_index
+        # if met guaranteed steps, look for next keyframe in case need to switch
+        if self._current_used_steps >= self._current_keyframe.guarantee_steps:
+            # if has next index, loop through and see if need t oswitch
+            if self.has_index(self._current_index+1):
+                for i in range(self._current_index+1, len(self.keyframes)):
+                    eval_c = self.keyframes[i]
+                    # check if start_t is greater or equal to curr_t
+                    # NOTE: t is in terms of sigmas, not percent, so bigger number = earlier step in sampling
+                    if eval_c.start_t >= curr_t:
+                        self._current_index = i
+                        self._current_keyframe = eval_c
+                        self._current_used_steps = 0
+                        # if guarantee_steps greater than zero, stop searching for other keyframes
+                        if self._current_keyframe.guarantee_steps > 0:
+                            break
+                    # if eval_c is outside the percent range, stop looking further
+                    else: break
+        # update steps current context is used
+        self._current_used_steps += 1
+
+    def patch_model(self, model: ModelPatcher) -> ModelPatcher:
+        def evolved_custom_cfg(args):
+            cond: Tensor = args["cond"]
+            uncond: Tensor = args["uncond"]
+            # cond scale is based purely off of CustomCFG - cond_scale input in sampler is ignored!
+            cond_scale = self.cfg_multival
+            if isinstance(cond_scale, Tensor):
+                cond_scale = prepare_mask_batch(cond_scale.to(cond.dtype).to(cond.device), cond.shape)
+                cond_scale = extend_to_batch_size(cond_scale, cond.shape[0])
+            return uncond + (cond - uncond) * cond_scale
+
+        model = model.clone()
+        model.set_model_sampler_cfg_function(evolved_custom_cfg)
+        return model
+
+    # properties shadow those of CustomCFGKeyframe
+    @property
+    def cfg_multival(self):
+        if self._current_keyframe != None:
+            return self._current_keyframe.cfg_multival
+        return None
@@ -16,7 +16,7 @@ from comfy.controlnet import ControlBase
 import comfy.ops

 from .context import ContextFuseMethod, ContextSchedules, get_context_weights, get_context_windows
-from .sample_settings import IterationOptions, SeedNoiseGeneration, prepare_mask_ad
+from .sample_settings import IterationOptions, SampleSettings, SeedNoiseGeneration, prepare_mask_ad
 from .utils_model import ModelTypeSD, wrap_function_to_inject_xformers_bug_info
 from .model_injection import InjectionParams, ModelPatcherAndInjector, MotionModelGroup, MotionModelPatcher
 from .motion_module_ad import AnimateDiffFormat, AnimateDiffInfo, AnimateDiffVersion, VanillaTemporalModule
@@ -30,6 +30,7 @@ class AnimateDiffHelper_GlobalState:
    def __init__(self):
        self.motion_models: MotionModelGroup = None
        self.params: InjectionParams = None
+        self.sample_settings: SampleSettings = None
        self.reset()
    
    def initialize(self, model):
@@ -40,6 +41,8 @@ class AnimateDiffHelper_GlobalState:
                self.motion_models.initialize_timesteps(model)
            if self.params.context_options is not None:
                self.params.context_options.initialize_timesteps(model)
+            if self.sample_settings.custom_cfg is not None:
+                self.sample_settings.custom_cfg.initialize_timesteps(model)

    def reset(self):
        self.initialized = False
@@ -53,6 +56,9 @@ class AnimateDiffHelper_GlobalState:
        if self.params is not None:
            del self.params
            self.params = None
+        if self.sample_settings is not None:
+            del self.sample_settings
+            self.sample_settings = None
    
    def update_with_inject_params(self, params: InjectionParams):
        self.params = params
@@ -147,8 +153,9 @@ def get_additional_models_factory(orig_get_additional_models: Callable, motion_m

 def apply_params_to_motion_models(motion_models: MotionModelGroup, params: InjectionParams):
    params = params.clone()
-    if params.context_options.context_schedule == ContextSchedules.VIEW_AS_CONTEXT:
-        params.context_options._current_context.context_length = params.full_length
+    for context in params.context_options.contexts:
+        if context.context_schedule == ContextSchedules.VIEW_AS_CONTEXT:
+            context.context_length = params.full_length
    # TODO: check (and message) should be different based on use_on_equal_length setting
    if params.context_options.context_length:
        pass
@@ -167,7 +174,7 @@ def apply_params_to_motion_models(motion_models: MotionModelGroup, params: Injec
        # if no context_length, treat video length as intended AD frame window
        if not params.context_options.context_length:
            for motion_model in motion_models.models:
-                if params.full_length > motion_model.model.encoding_max_len:
+                if not motion_model.model.is_length_valid_for_encoding_max_len(params.full_length):
                    raise ValueError(f"Without a context window, AnimateDiff model {motion_model.model.mm_info.mm_name} has upper limit of {motion_model.model.encoding_max_len} frames, but received {params.full_length} latents.")
            motion_models.set_video_length(params.full_length, params.full_length)
        # otherwise, treat context_length as intended AD frame window
@@ -175,7 +182,7 @@ def apply_params_to_motion_models(motion_models: MotionModelGroup, params: Injec
            for motion_model in motion_models.models:
                view_options = params.context_options.view_options
                context_length = view_options.context_length if view_options else params.context_options.context_length
-                if context_length > motion_model.model.encoding_max_len:
+                if not motion_model.model.is_length_valid_for_encoding_max_len(context_length):
                    raise ValueError(f"AnimateDiff model {motion_model.model.mm_info.mm_name} has upper limit of {motion_model.model.encoding_max_len} frames for a context window, but received context length of {params.context_options.context_length}.")
            motion_models.set_video_length(params.context_options.context_length, params.full_length)
        # inject model
@@ -202,10 +209,10 @@ class FunctionInjectionHolder:
        if params.unlimited_area_hack:
            model.model.memory_required = unlimited_memory_required
        if model.motion_models is not None:
-            # only apply groupnorm hack if not [v3 or (AnimateDiff SD1.5 and v2 and should apply v2 properly)]
+            # only apply groupnorm hack if not [v3 or ([not Hotshot] and SD1.5 and v2 and apply_v2_properly)]
            info: AnimateDiffInfo = model.motion_models[0].model.mm_info
-            if not (info.mm_version == AnimateDiffVersion.V3 or (info.mm_format == AnimateDiffFormat.ANIMATEDIFF and info.sd_type == ModelTypeSD.SD1_5 and
-                    info.mm_version == AnimateDiffVersion.V2 and params.apply_v2_properly)):
+            if not (info.mm_version == AnimateDiffVersion.V3 or
+                    (info.mm_format not in [AnimateDiffFormat.HOTSHOTXL] and info.sd_type == ModelTypeSD.SD1_5 and info.mm_version == AnimateDiffVersion.V2 and params.apply_v2_properly)):
                torch.nn.GroupNorm.forward = groupnorm_mm_factory(params)
                comfy.ops.manual_cast.GroupNorm.forward_comfy_cast_weights = groupnorm_mm_factory(params, manual_cast=True)
                # if mps device (Apple Silicon), disable batched conds to avoid black images with groupnorm hack
@@ -245,6 +252,8 @@ def motion_sample_factory(orig_comfy_sample: Callable, is_custom: bool=False) ->
        cached_noise = None
        function_injections = FunctionInjectionHolder()
        try:
+            if model.sample_settings.custom_cfg is not None:
+                model = model.sample_settings.custom_cfg.patch_model(model)
            # clone params from model
            params = model.motion_injection_params.clone()
            # get amount of latents passed in, and store in params
@@ -279,6 +288,7 @@ def motion_sample_factory(orig_comfy_sample: Callable, is_custom: bool=False) ->
                ADGS.current_step = ADGS.start_step + step + 1
            kwargs["callback"] = ad_callback
            ADGS.motion_models = model.motion_models
+            ADGS.sample_settings = model.sample_settings

            # apply adapt_denoise_steps
            args = list(args)
@@ -331,6 +341,8 @@ def motion_sample_factory(orig_comfy_sample: Callable, is_custom: bool=False) ->

                if model.motion_models is not None:
                    model.motion_models.pre_run(model)
+                if model.sample_settings is not None:
+                    model.sample_settings.pre_run(model)
                latents = wrap_function_to_inject_xformers_bug_info(orig_comfy_sample)(model, noise, *args, **kwargs)
            return latents
        finally:
@@ -352,8 +364,11 @@ def evolved_sampling_function(model, x, timestep, uncond, cond, cond_scale, mode
        ADGS.motion_models.prepare_current_keyframe(t=timestep)
    if ADGS.params.context_options is not None:
        ADGS.params.context_options.prepare_current_context(t=timestep)
+    if ADGS.sample_settings.custom_cfg is not None:
+        ADGS.sample_settings.custom_cfg.prepare_current_keyframe(t=timestep)

-    if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
+    # never use cfg1 optimization if using custom_cfg (since can have timesteps and such)
+    if ADGS.sample_settings.custom_cfg is None and math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
        uncond_ = None
    else:
        uncond_ = uncond
@@ -454,6 +469,7 @@ def sliding_calc_cond_uncond_batch(model, cond, uncond, x_in: Tensor, timestep,
    cond_final = torch.zeros_like(x_in)
    uncond_final = torch.zeros_like(x_in)
    out_count_final = torch.zeros((x_in.shape[0], 1, 1, 1), device=x_in.device)
+    bias_final = [0.0] * x_in.shape[0]

    # perform calc_cond_uncond_batch per context window
    for ctx_idxs in context_windows:
@@ -477,16 +493,36 @@ def sliding_calc_cond_uncond_batch(model, cond, uncond, x_in: Tensor, timestep,

        sub_cond_out, sub_uncond_out = comfy.samplers.calc_cond_uncond_batch(model, sub_cond, sub_uncond, sub_x, sub_timestep, model_options)

-        # add conds and counts based on weights of fuse method
-        weights = get_context_weights(len(ctx_idxs), ADGS.params.context_options.fuse_method) * batched_conds
-        weights_tensor = torch.Tensor(weights).to(device=x_in.device).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
-        cond_final[full_idxs] += sub_cond_out * weights_tensor
-        uncond_final[full_idxs] += sub_uncond_out * weights_tensor
-        out_count_final[full_idxs] += weights_tensor
+        if ADGS.params.context_options.fuse_method == ContextFuseMethod.RELATIVE:
+            full_length = ADGS.params.full_length
+            for pos, idx in enumerate(ctx_idxs):
+                # bias is the influence of a specific index in relation to the whole context window
+                bias = 1 - abs(idx - (ctx_idxs[0] + ctx_idxs[-1]) / 2) / ((ctx_idxs[-1] - ctx_idxs[0] + 1e-2) / 2)
+                bias = max(1e-2, bias)
+                # take weighted average relative to total bias of current idx
+                # and account for batched_conds
+                for n in range(batched_conds):
+                    bias_total = bias_final[(full_length*n)+idx]
+                    prev_weight = (bias_total / (bias_total + bias))
+                    new_weight = (bias / (bias_total + bias))
+                    cond_final[(full_length*n)+idx] = cond_final[(full_length*n)+idx] * prev_weight + sub_cond_out[(full_length*n)+pos] * new_weight
+                    uncond_final[(full_length*n)+idx] = uncond_final[(full_length*n)+idx] * prev_weight + sub_uncond_out[(full_length*n)+pos] * new_weight
+                    bias_final[(full_length*n)+idx] = bias_total + bias
+        else:
+            # add conds and counts based on weights of fuse method
+            weights = get_context_weights(len(ctx_idxs), ADGS.params.context_options.fuse_method) * batched_conds
+            weights_tensor = torch.Tensor(weights).to(device=x_in.device).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+            cond_final[full_idxs] += sub_cond_out * weights_tensor
+            uncond_final[full_idxs] += sub_uncond_out * weights_tensor
+            out_count_final[full_idxs] += weights_tensor

-
-    # normalize cond and uncond via division by context usage counts
-    cond_final /= out_count_final
-    uncond_final /= out_count_final
-    del out_count_final
-    return cond_final, uncond_final
+    if ADGS.params.context_options.fuse_method == ContextFuseMethod.RELATIVE:
+        # already normalized, so return as is
+        del out_count_final
+        return cond_final, uncond_final
+    else:
+        # normalize cond and uncond via division by context usage counts
+        cond_final /= out_count_final
+        uncond_final /= out_count_final
+        del out_count_final
+        return cond_final, uncond_final
@@ -3,72 +3,169 @@ from pathlib import Path
 from typing import Callable, Union
 from collections.abc import Iterable
 from time import time
+import copy

+import torch
 import numpy as np

 import folder_paths
-from comfy.model_base import SD21UNCLIP, SDXL, BaseModel, SDXLRefiner, SVD_img2vid, model_sampling
+from comfy.model_base import SD21UNCLIP, SDXL, BaseModel, SDXLRefiner, SVD_img2vid, model_sampling, ModelType
 from comfy.model_management import xformers_enabled
 from comfy.model_patcher import ModelPatcher

+import comfy.model_sampling
+import comfy_extras.nodes_model_advanced

-class IsChangedHelper:
-    def __init__(self):
-        self.val = 0
-    
-    def no_change(self):
-        return self.val
-    
-    def change(self):
-        self.val = (self.val + 1) % 100
+
+BIGMIN = -(2**53-1)
+BIGMAX = (2**53-1)


 class ModelSamplingConfig:
-    def __init__(self, beta_schedule: str):
+    def __init__(self, beta_schedule: str, linear_start: float=None, linear_end: float=None):
        self.sampling_settings = {"beta_schedule": beta_schedule}
+        if linear_start is not None:
+            self.sampling_settings["linear_start"] = linear_start
+        if linear_end is not None:
+            self.sampling_settings["linear_end"] = linear_end
        self.beta_schedule = beta_schedule  # keeping this for backwards compatibility


+class ModelSamplingType:
+    EPS = "eps"
+    V_PREDICTION = "v_prediction"
+    LCM = "lcm"
+
+    _NON_LCM_LIST = [EPS, V_PREDICTION]
+    _FULL_LIST = [EPS, V_PREDICTION, LCM]
+
+    MAP = {
+        EPS: ModelType.EPS,
+        V_PREDICTION: ModelType.V_PREDICTION,
+        LCM: comfy_extras.nodes_model_advanced.LCM,
+    }
+
+    @classmethod
+    def from_alias(cls, alias: str):
+        return cls.MAP[alias]
+
+
+def factory_model_sampling_discrete_distilled(original_timesteps=50):
+    class ModelSamplingDiscreteDistilledEvolved(comfy_extras.nodes_model_advanced.ModelSamplingDiscreteDistilled):
+        def __init__(self, *args, **kwargs):
+            self.original_timesteps = original_timesteps  # normal LCM has 50
+            super().__init__(*args, **kwargs)
+    return ModelSamplingDiscreteDistilledEvolved
+
+
+# based on code in comfy_extras/nodes_model_advanced.py
+def evolved_model_sampling(model_config: ModelSamplingConfig, model_type: ModelType, alias: str, original_timesteps: int=None):
+    # if LCM, need to handle manually
+    if BetaSchedules.is_lcm(alias) or original_timesteps is not None:
+        sampling_type = comfy_extras.nodes_model_advanced.LCM
+        if original_timesteps is not None:
+            sampling_base = factory_model_sampling_discrete_distilled(original_timesteps=original_timesteps)
+        elif alias == BetaSchedules.LCM_100:
+            sampling_base = factory_model_sampling_discrete_distilled(original_timesteps=100)
+        elif alias == BetaSchedules.LCM_25:
+            sampling_base = factory_model_sampling_discrete_distilled(original_timesteps=25)
+        else:
+            sampling_base = comfy_extras.nodes_model_advanced.ModelSamplingDiscreteDistilled
+        class ModelSamplingAdvancedEvolved(sampling_base, sampling_type):
+            pass
+        # NOTE: if I want to support zsnr, this is where I would add that code
+        return ModelSamplingAdvancedEvolved(model_config)
+    # otherwise, use vanilla model_sampling function
+    return model_sampling(model_config, model_type)
+
+
 class BetaSchedules:
    AUTOSELECT = "autoselect"
    SQRT_LINEAR = "sqrt_linear (AnimateDiff)"
    LINEAR_ADXL = "linear (AnimateDiff-SDXL)"
    LINEAR = "linear (HotshotXL/default)"
+    AVG_LINEAR_SQRT_LINEAR = "avg(sqrt_linear,linear)"
+    LCM_AVG_LINEAR_SQRT_LINEAR = "lcm avg(sqrt_linear,linear)"
+    LCM = "lcm"
+    LCM_100 = "lcm[100_ots]"
+    LCM_25 = "lcm[25_ots]"
+    LCM_SQRT_LINEAR = "lcm >> sqrt_linear"
    USE_EXISTING = "use existing"
    SQRT = "sqrt"
    COSINE = "cosine"
    SQUAREDCOS_CAP_V2 = "squaredcos_cap_v2"
+    RAW_LINEAR = "linear"
+    RAW_SQRT_LINEAR = "sqrt_linear"

-    ALIAS_LIST = [AUTOSELECT, SQRT_LINEAR, LINEAR_ADXL, LINEAR,
-                  USE_EXISTING, SQRT, COSINE, SQUAREDCOS_CAP_V2]
+    RAW_BETA_SCHEDULE_LIST = [RAW_LINEAR, RAW_SQRT_LINEAR, SQRT, COSINE, SQUAREDCOS_CAP_V2]
+
+    ALIAS_LCM_LIST = [LCM, LCM_100, LCM_25, LCM_SQRT_LINEAR]
+
+    ALIAS_ACTIVE_LIST = [SQRT_LINEAR, LINEAR_ADXL, LINEAR, AVG_LINEAR_SQRT_LINEAR, LCM_AVG_LINEAR_SQRT_LINEAR, LCM, LCM_100, LCM_SQRT_LINEAR, # LCM_25 is purposely omitted
+                  SQRT, COSINE, SQUAREDCOS_CAP_V2]
+
+    ALIAS_LIST = [AUTOSELECT, USE_EXISTING] + ALIAS_ACTIVE_LIST
+
+    

    ALIAS_MAP = {
        SQRT_LINEAR: "sqrt_linear",
        LINEAR_ADXL: "linear", # also linear, but has different linear_end (0.020)
        LINEAR: "linear",
+        LCM_100: "linear",  # distilled, 100 original timesteps
+        LCM_25: "linear",  # distilled, 25 original timesteps
+        LCM: "linear",  # distilled
+        LCM_SQRT_LINEAR: "sqrt_linear", # distilled, sqrt_linear
        SQRT: "sqrt",
        COSINE: "cosine",
        SQUAREDCOS_CAP_V2: "squaredcos_cap_v2",
+        RAW_LINEAR: "linear",
+        RAW_SQRT_LINEAR: "sqrt_linear"
    }

+    @classmethod
+    def is_lcm(cls, alias: str):
+        return alias in cls.ALIAS_LCM_LIST
+
    @classmethod
    def to_name(cls, alias: str):
        return cls.ALIAS_MAP[alias]
    
    @classmethod
    def to_config(cls, alias: str) -> ModelSamplingConfig:
-        return ModelSamplingConfig(cls.to_name(alias))
-    
-    @classmethod
-    def to_model_sampling(cls, alias: str, model: ModelPatcher):
-        if alias == cls.USE_EXISTING:
-            return None
-        ms_obj = model_sampling(cls.to_config(alias), model_type=model.model.model_type)
+        linear_start = None
+        linear_end = None
        if alias == cls.LINEAR_ADXL:
            # uses linear_end=0.020
-            ms_obj._register_schedule(given_betas=None, beta_schedule=cls.to_name(alias), timesteps=1000, linear_start=0.00085, linear_end=0.020, cosine_s=8e-3)
+            linear_end = 0.020
+        return ModelSamplingConfig(cls.to_name(alias), linear_start=linear_start, linear_end=linear_end)
+    
+    @classmethod
+    def _to_model_sampling(cls, alias: str, model_type: ModelType, config_override: ModelSamplingConfig=None, original_timesteps: int=None):
+        if alias == cls.USE_EXISTING:
+            return None
+        elif config_override != None:
+            return evolved_model_sampling(config_override, model_type=model_type, alias=alias, original_timesteps=original_timesteps)
+        elif alias == cls.AVG_LINEAR_SQRT_LINEAR:
+            ms_linear = evolved_model_sampling(cls.to_config(cls.LINEAR), model_type=model_type, alias=cls.LINEAR)
+            ms_sqrt_linear = evolved_model_sampling(cls.to_config(cls.SQRT_LINEAR), model_type=model_type, alias=cls.SQRT_LINEAR)
+            avg_sigmas = (ms_linear.sigmas + ms_sqrt_linear.sigmas) / 2
+            ms_linear.set_sigmas(avg_sigmas)
+            return ms_linear
+        elif alias == cls.LCM_AVG_LINEAR_SQRT_LINEAR:
+            ms_linear = evolved_model_sampling(cls.to_config(cls.LCM), model_type=model_type, alias=cls.LCM)
+            ms_sqrt_linear = evolved_model_sampling(cls.to_config(cls.LCM_SQRT_LINEAR), model_type=model_type, alias=cls.LCM_SQRT_LINEAR)
+            avg_sigmas = (ms_linear.sigmas + ms_sqrt_linear.sigmas) / 2
+            ms_linear.set_sigmas(avg_sigmas)
+            return ms_linear
+            # average out the sigmas
+        ms_obj = evolved_model_sampling(cls.to_config(alias), model_type=model_type, alias=alias, original_timesteps=original_timesteps)
        return ms_obj

+    @classmethod
+    def to_model_sampling(cls, alias: str, model: ModelPatcher):
+        return cls._to_model_sampling(alias=alias, model_type=model.model.model_type)
+
    @staticmethod
    def get_alias_list_with_first_element(first_element: str):
        new_list = BetaSchedules.ALIAS_LIST.copy()
@@ -77,16 +174,65 @@ class BetaSchedules:
        return new_list


-class BetaScheduleCache:
-    def __init__(self, model: ModelPatcher): 
-        self.model_sampling = model.model.model_sampling
+class SigmaSchedule:
+    def __init__(self, model_sampling: comfy.model_sampling.ModelSamplingDiscrete, model_type: ModelType):
+        self.model_sampling = model_sampling
+        #self.config = config
+        self.model_type = model_type
+        self.original_timesteps = getattr(self.model_sampling, "original_timesteps", None)
+    
+    def is_lcm(self):
+        return self.original_timesteps is not None

-    def use_cached_beta_schedule_and_clean(self, model: ModelPatcher):
-        model.model.model_sampling = self.model_sampling
-        self.clean()
+    def total_sigmas(self):
+        return len(self.model_sampling.sigmas)
+    
+    def clone(self) -> 'SigmaSchedule':
+        new_model_sampling = copy.deepcopy(self.model_sampling)
+        #new_config = copy.deepcopy(self.config)
+        return SigmaSchedule(model_sampling=new_model_sampling, model_type=self.model_type)

-    def clean(self):
-        self.model_sampling = None
+    # def clone(self):
+    #     pass
+
+    @staticmethod
+    def apply_zsnr(new_model_sampling: comfy.model_sampling.ModelSamplingDiscrete):
+        new_model_sampling.set_sigmas(comfy_extras.nodes_model_advanced.rescale_zero_terminal_snr_sigmas(new_model_sampling.sigmas))
+
+    # def get_lcmified(self, original_timesteps=50, zsnr=False) -> 'SigmaSchedule':
+    #     new_model_sampling = evolved_model_sampling(model_config=self.config, model_type=self.model_type, alias=None, original_timesteps=original_timesteps)
+    #     if zsnr:
+    #         new_model_sampling.set_sigmas(comfy_extras.nodes_model_advanced.rescale_zero_terminal_snr_sigmas(new_model_sampling.sigmas))
+    #     return SigmaSchedule(model_sampling=new_model_sampling, config=self.config, model_type=self.model_type, is_lcm=True)
+        
+
+class InterpolationMethod:
+    LINEAR = "linear"
+    EASE_IN = "ease_in"
+    EASE_OUT = "ease_out"
+    EASE_IN_OUT = "ease_in_out"
+
+    _LIST = [LINEAR, EASE_IN, EASE_OUT, EASE_IN_OUT]
+
+    @classmethod
+    def get_weights(cls, num_from: float, num_to: float, length: int, method: str, reverse=False):
+        diff = num_to - num_from
+        if method == cls.LINEAR:
+            weights = torch.linspace(num_from, num_to, length)
+        elif method == cls.EASE_IN:
+            index = torch.linspace(0, 1, length)
+            weights = diff * np.power(index, 2) + num_from
+        elif method == cls.EASE_OUT:
+            index = torch.linspace(0, 1, length)
+            weights = diff * (1 - np.power(1 - index, 2)) + num_from
+        elif method == cls.EASE_IN_OUT:
+            index = torch.linspace(0, 1, length)
+            weights = diff * ((1 - np.cos(index * np.pi)) / 2) + num_from
+        else:
+            raise ValueError(f"Unrecognized interpolation method '{method}'.")
+        if reverse:
+            weights = weights.flip(dims=(0,))
+        return weights


 class Folders:
@@ -167,8 +313,6 @@ def calculate_model_hash(model: ModelPatcher):
        m.update(buf.cpu().numpy().view(np.uint8))
    return m.hexdigest()

-BIGMIN = -(2**63-1)
-BIGMAX = (2**63-1)

 class ModelTypeSD:
    SD1_5 = "SD1.5"