Bugfix encode step for XY_Capsule (inspire pack) (#187)

* bugfix: in case of XY_Capsule, prompt needs to be encoded by the analysed LoRA's clip

* bugfix: when processing LoRAs you need to reload the LoRA every index on the y scale (not only the first one) to maintain the LoRA-patched clip.

* add align your step schedulers via monkey patching

---------

Co-authored-by: larsupb <larsupb@gmx.net>

efficiency_nodes.py

@@ -18,6 +18,8 @@ import subprocess
 import json
 import psutil
 
+from comfy_extras.nodes_align_your_steps import AlignYourStepsScheduler
+
 # Get the absolute path of various directories
 my_dir = os.path.dirname(os.path.abspath(__file__))
 custom_nodes_dir = os.path.abspath(os.path.join(my_dir, '..'))
@@ -52,12 +54,17 @@ from .py import bnk_tiled_samplers
 from .py import bnk_adv_encode
 sys.path.remove(my_dir)
 
+from comfy import samplers
 # Append custom_nodes_dir to sys.path
 sys.path.append(custom_nodes_dir)
 
 # GLOBALS
 REFINER_CFG_OFFSET = 0 #Refiner CFG Offset
 
+# Monkey patch schedulers
+samplers.SCHEDULER_NAMES = samplers.SCHEDULER_NAMES + ["AYS SD1", "AYS SDXL", "AYS SVD"]
+samplers.KSampler.SCHEDULERS = samplers.KSampler.SCHEDULERS + ["AYS SD1", "AYS SDXL", "AYS SVD"]
+
 ########################################################################################################################
 # Common function for encoding prompts
 def encode_prompts(positive_prompt, negative_prompt, token_normalization, weight_interpretation, clip, clip_skip,
@@ -394,7 +401,6 @@ class TSC_Apply_ControlNet_Stack:
 ########################################################################################################################
 # TSC KSampler (Efficient)
 class TSC_KSampler:
-    
     empty_image = pil2tensor(Image.new('RGBA', (1, 1), (0, 0, 0, 0)))
 
     @classmethod
@@ -429,6 +435,15 @@ class TSC_KSampler:
                optional_vae=(None,), script=None, add_noise=None, start_at_step=None, end_at_step=None,
                return_with_leftover_noise=None, sampler_type="regular"):
 
+        # monkey patch the sample function
+        original_calculation = comfy.samplers.calculate_sigmas
+        def calculate_sigmas(model_sampling, scheduler_name: str, steps):
+            if scheduler_name.startswith("AYS"):
+                return AlignYourStepsScheduler().get_sigmas(scheduler_name.split(" ")[1], steps, denoise=1.0)[0]
+            return original_calculation(model_sampling, scheduler_name, steps)
+        comfy.samplers.SCHEDULER_NAMES = comfy.samplers.SCHEDULER_NAMES + ["AYS SD1", "AYS SDXL", "AYS SVD"]
+        comfy.samplers.calculate_sigmas = calculate_sigmas
+
         # Rename the vae variable
         vae = optional_vae
 
@@ -1363,7 +1378,7 @@ class TSC_KSampler:
                     encode = True
 
                 # Load LoRA if required
-                elif (X_type == "LoRA" and index == 0):
+                elif (X_type == "LoRA"):
                     # Don't cache Checkpoints
                     model, clip = load_lora(lora_stack, ckpt_name, xyplot_id, cache=cache[2])
                     encode = True
@@ -1383,13 +1398,13 @@ class TSC_KSampler:
 
                 # Encode base prompt if required
                 encode_types = ["Positive Prompt S/R", "Negative Prompt S/R", "Clip Skip", "ControlNetStrength",
-                                "ControlNetStart%",  "ControlNetEnd%"]
+                                "ControlNetStart%",  "ControlNetEnd%", "XY_Capsule"]
                 if (X_type in encode_types and index == 0) or Y_type in encode_types:
                     encode = True
 
                 # Encode refiner prompt if required
                 encode_refiner_types = ["Positive Prompt S/R", "Negative Prompt S/R", "AScore+", "AScore-",
-                                        "Clip Skip (Refiner)"]
+                                        "Clip Skip (Refiner)", "XY_Capsule"]
                 if (X_type in encode_refiner_types and index == 0) or Y_type in encode_refiner_types:
                     encode_refiner = True