diff --git a/efficiency_nodes.py b/efficiency_nodes.py index 076d758..2af4201 100644 --- a/efficiency_nodes.py +++ b/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