feat: Samplers - support AYS scheduler

fix: invalid step bounding of KSamplerAdvanced_progress !! requirements: Impact Pack V4.90 or later
2024-04-27 16:41:13 +09:00
parent e9de5befd6
commit 92cd5108c4
5 changed files with 31 additions and 38 deletions
@@ -1,11 +1,10 @@
 import comfy
 import torch
-import numpy as np
-import latent_preview
 from .libs import utils
 from einops import rearrange
 import random
 import math
+from .libs import common


 def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0,
@@ -23,29 +22,8 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
            batch_inds = latent["batch_index"] if "batch_index" in latent else None
            noise = utils.prepare_noise(latent_image, seed, batch_inds, noise_device, incremental_seed_mode, variation_seed=variation_seed, variation_strength=variation_strength)

-    noise_mask = None
-    if "noise_mask" in latent:
-        noise_mask = latent["noise_mask"]
-
-    preview_format = "JPEG"
-    if preview_format not in ["JPEG", "PNG"]:
-        preview_format = "JPEG"
-
-    previewer = latent_preview.get_previewer(device, model.model.latent_format)
-
-    pbar = comfy.utils.ProgressBar(steps)
-    def callback(step, x0, x, total_steps):
-        preview_bytes = None
-        if previewer:
-            preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
-        pbar.update_absolute(step + 1, total_steps, preview_bytes)
-
-    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
-                                  denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step,
-                                  force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, seed=seed)
-    out = latent.copy()
-    out["samples"] = samples
-    return (out, )
+    samples = common.impact_sampling(model, not disable_noise, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, start_step, last_step, not force_full_denoise, noise=noise)
+    return (samples, )


 class KSampler_inspire:
@@ -57,7 +35,7 @@ class KSampler_inspire:
                     "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                     "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
                     "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                     "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                     "scheduler": (common.SCHEDULERS, ),
                     "positive": ("CONDITIONING", ),
                     "negative": ("CONDITIONING", ),
                     "latent_image": ("LATENT", ),
@@ -88,7 +66,7 @@ class KSamplerAdvanced_inspire:
                     "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                     "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}),
                     "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                     "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                     "scheduler": (common.SCHEDULERS, ),
                     "positive": ("CONDITIONING", ),
                     "negative": ("CONDITIONING", ),
                     "latent_image": ("LATENT", ),
@@ -137,7 +115,7 @@ class KSampler_inspire_pipe:
                     "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                     "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
                     "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                     "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                     "scheduler": (common.SCHEDULERS, ),
                     "latent_image": ("LATENT", ),
                     "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                     "noise_mode": (["GPU(=A1111)", "CPU"],),
@@ -168,7 +146,7 @@ class KSamplerAdvanced_inspire_pipe:
                     "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                     "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}),
                     "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                     "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                     "scheduler": (common.SCHEDULERS, ),
                     "latent_image": ("LATENT", ),
                     "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
                     "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
@@ -0,0 +1,14 @@
+import comfy
+import nodes
+from . import utils
+
+SCHEDULERS = comfy.samplers.KSampler.SCHEDULERS + ['AYS SDXL', 'AYS SD1', 'AYS SVD']
+
+
+def impact_sampling(*args, **kwargs):
+    if 'RegionalSampler' not in nodes.NODE_CLASS_MAPPINGS:
+        utils.try_install_custom_node('https://github.com/ltdrdata/ComfyUI-Impact-Pack',
+                                      "'Impact Pack' extension is required.")
+        raise Exception(f"[ERROR] You need to install 'ComfyUI-Impact-Pack'")
+
+    return nodes.NODE_CLASS_MAPPINGS['RegionalSampler'].separated_sample(*args, **kwargs)
@@ -15,7 +15,7 @@ import traceback
 import random

 from server import PromptServer
-from .libs import utils
+from .libs import utils, common
 from .backend_support import CheckpointLoaderSimpleShared

 prompt_builder_preset = {}
@@ -344,7 +344,7 @@ class GlobalSampler:
        return {
            "required": {
                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                "scheduler": (common.SCHEDULERS, ),
            }
        }

@@ -4,7 +4,7 @@ import comfy
 import nodes
 import torch
 from . import prompt_support
-from .libs import utils
+from .libs import utils, common


 class RegionalPromptSimple:
@@ -16,7 +16,7 @@ class RegionalPromptSimple:
                "mask": ("MASK",),
                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
                "sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
-                "scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
+                "scheduler": (common.SCHEDULERS,),
                "wildcard_prompt": ("STRING", {"multiline": True, "dynamicPrompts": False, "placeholder": "wildcard prompt"}),
                "controlnet_in_pipe": ("BOOLEAN", {"default": False, "label_on": "Keep", "label_off": "Override"}),
                "sigma_factor": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
@@ -91,7 +91,7 @@ class RegionalPromptColorMask:
                "mask_color": ("STRING", {"multiline": False, "default": "#FFFFFF"}),
                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
                "sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
-                "scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
+                "scheduler": (common.SCHEDULERS,),
                "wildcard_prompt": ("STRING", {"multiline": True, "dynamicPrompts": False, "placeholder": "wildcard prompt"}),
                "controlnet_in_pipe": ("BOOLEAN", {"default": False, "label_on": "Keep", "label_off": "Override"}),
                "sigma_factor": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
@@ -1,6 +1,8 @@
 import torch
 from . import a1111_compat
 import comfy
+from .libs import common
+

 class KSampler_progress(a1111_compat.KSampler_inspire):
    @classmethod
@@ -11,7 +13,7 @@ class KSampler_progress(a1111_compat.KSampler_inspire):
                     "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                     "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
                     "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                     "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                     "scheduler": (common.SCHEDULERS, ),
                     "positive": ("CONDITIONING", ),
                     "negative": ("CONDITIONING", ),
                     "latent_image": ("LATENT", ),
@@ -63,7 +65,7 @@ class KSamplerAdvanced_progress(a1111_compat.KSamplerAdvanced_inspire):
                     "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                     "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}),
                     "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                     "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                     "scheduler": (common.SCHEDULERS, ),
                     "positive": ("CONDITIONING", ),
                     "negative": ("CONDITIONING", ),
                     "latent_image": ("LATENT", ),
@@ -92,11 +94,10 @@ class KSamplerAdvanced_progress(a1111_compat.KSamplerAdvanced_inspire):
        else:
            result = [latent_image['samples']]

-        for i in range(start_at_step, end_at_step+1):
+        for i in range(start_at_step, min(end_at_step+1, steps+1)):
            cur_add_noise = i == start_at_step and add_noise
            cur_return_with_leftover_noise = i != steps or return_with_leftover_noise
            latent_image = sampler.sample(model, cur_add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, i, i+1, noise_mode, cur_return_with_leftover_noise)[0]
-            print(f"{i}, {i+1}")
            if i % interval == 0 or i == steps:
                result.append(latent_image['samples'])