feat: ImpactKSamplers supports AYS scheduler

refactor: impact_sampling
- add: sample_with_custom_noise
- expose separated_sample api through RegionalSampler
This commit is contained in:
Dr.Lt.Data
2024-04-27 16:39:07 +09:00
parent f283125334
commit f7532d8bce
4 changed files with 78 additions and 16 deletions
+1 -1
View File
@@ -2,7 +2,7 @@ import configparser
import os
version_code = [4, 89]
version_code = [4, 90]
version = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' if len(version_code) > 2 else '')
dependency_version = 20
-4
View File
@@ -2150,7 +2150,3 @@ class ImpactWildcardEncode:
populated = kwargs['populated_text']
model, clip, conditioning = impact.wildcards.process_with_loras(populated, kwargs['model'], kwargs['clip'])
return (model, clip, conditioning, populated)
+66 -4
View File
@@ -15,7 +15,7 @@ def calculate_sigmas(model, sampler, scheduler, steps):
if hasattr(samplers, 'calculate_sigmas'):
if scheduler.startswith('AYS'):
sigmas = nodes.NODE_CLASS_MAPPINGS['AlignYourStepsScheduler']().get_sigmas(scheduler[4:], steps)[0]
sigmas = nodes.NODE_CLASS_MAPPINGS['AlignYourStepsScheduler']().get_sigmas(scheduler[4:], steps, denoise=1.0)[0]
else:
sigmas = samplers.calculate_sigmas(model.get_model_object("model_sampling"), scheduler, steps)
else:
@@ -101,14 +101,68 @@ def ksampler(sampler_name, total_sigmas, extra_options={}, inpaint_options={}):
return samplers.KSAMPLER(sampler_function, extra_options, inpaint_options)
from comfy_extras.nodes_custom_sampler import Noise_EmptyNoise, Noise_RandomNoise
import latent_preview
import comfy
# modified version of SamplerCustom.sample
def sample_with_custom_noise(model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image, noise=None):
latent = latent_image
latent_image = latent["samples"]
if noise is None:
if not add_noise:
noise = Noise_EmptyNoise().generate_noise(latent)
else:
noise = Noise_RandomNoise(noise_seed).generate_noise(latent)
noise_mask = None
if "noise_mask" in latent:
noise_mask = latent["noise_mask"]
x0_output = {}
callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
out = latent.copy()
out["samples"] = samples
if "x0" in x0_output:
out_denoised = latent.copy()
out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
else:
out_denoised = out
return out, out_denoised
# When sampling one step at a time, it mitigates the problem. (especially for _sde series samplers)
def separated_sample(model, add_noise, seed, steps, cfg, sampler_name, scheduler, positive, negative,
latent_image, start_at_step, end_at_step, return_with_leftover_noise, sigma_ratio=1.0, sampler_opt=None):
latent_image, start_at_step, end_at_step, return_with_leftover_noise, sigma_ratio=1.0, sampler_opt=None, noise=None):
# still cannot preserve for 2s, 2m, 3m. that requires some more thing.
if sampler_opt is None:
total_sigmas = calculate_sigmas(model, sampler_name, scheduler, steps)
else:
total_sigmas = calculate_sigmas(model, "", scheduler, steps)
sigmas = total_sigmas[start_at_step:end_at_step+1] * sigma_ratio
sigmas = total_sigmas
if end_at_step is not None and end_at_step < (len(total_sigmas) - 1):
sigmas = total_sigmas[:end_at_step + 1]
if not return_with_leftover_noise:
sigmas[-1] = 0
if start_at_step is not None:
if start_at_step < (len(sigmas) - 1):
sigmas = sigmas[start_at_step:] * sigma_ratio
else:
if latent_image is not None:
return latent_image
else:
return torch.zeros_like(noise)
if sampler_opt is None:
impact_sampler = ksampler(sampler_name, total_sigmas)
else:
@@ -117,7 +171,7 @@ def separated_sample(model, add_noise, seed, steps, cfg, sampler_name, scheduler
if len(sigmas) == 0 or (len(sigmas) == 1 and sigmas[0] == 0):
return latent_image
res = nodes_custom_sampler.SamplerCustom().sample(model, add_noise, seed, cfg, positive, negative, impact_sampler, sigmas, latent_image)
res = sample_with_custom_noise(model, add_noise, seed, cfg, positive, negative, impact_sampler, sigmas, latent_image, noise=noise)
if return_with_leftover_noise:
return res[0]
@@ -125,6 +179,14 @@ def separated_sample(model, add_noise, seed, steps, cfg, sampler_name, scheduler
return res[1]
def impact_sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, sigma_ratio=1.0, sampler_opt=None, noise=None):
advanced_steps = math.floor(steps / denoise)
start_at_step = advanced_steps - steps
end_at_step = start_at_step + steps
return separated_sample(model, True, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, False)
def ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise,
refiner_ratio=None, refiner_model=None, refiner_clip=None, refiner_positive=None, refiner_negative=None, sigma_factor=1.0):
+11 -7
View File
@@ -3,7 +3,7 @@ import impact.core as core
from impact.utils import *
from nodes import MAX_RESOLUTION
import nodes
from impact.impact_sampling import KSamplerWrapper, KSamplerAdvancedWrapper
from impact.impact_sampling import KSamplerWrapper, KSamplerAdvancedWrapper, separated_sample, impact_sample
class TiledKSamplerProvider:
@@ -43,7 +43,7 @@ class KSamplerProvider:
"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": (core.SCHEDULERS, ),
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"basic_pipe": ("BASIC_PIPE", )
},
@@ -66,7 +66,7 @@ class KSamplerAdvancedProvider:
return {"required": {
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
"scheduler": (core.SCHEDULERS, ),
"sigma_factor": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
"basic_pipe": ("BASIC_PIPE", )
},
@@ -302,6 +302,10 @@ class RegionalSampler:
CATEGORY = "ImpactPack/Regional"
@staticmethod
def separated_sample(*args, **kwargs):
return separated_sample(*args, **kwargs)
@staticmethod
def mask_erosion(samples, mask, grow_mask_by):
mask = mask.clone()
@@ -540,7 +544,7 @@ class KSamplerBasicPipe:
"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": (core.SCHEDULERS, ),
"latent_image": ("LATENT", ),
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
}
@@ -553,7 +557,7 @@ class KSamplerBasicPipe:
def sample(self, basic_pipe, seed, steps, cfg, sampler_name, scheduler, latent_image, denoise=1.0):
model, clip, vae, positive, negative = basic_pipe
latent = nodes.KSampler().sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise)[0]
latent = impact_sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise)
return basic_pipe, latent, vae
@@ -567,7 +571,7 @@ class KSamplerAdvancedBasicPipe:
"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": (core.SCHEDULERS, ),
"latent_image": ("LATENT", ),
"start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
"end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
@@ -593,6 +597,6 @@ class KSamplerAdvancedBasicPipe:
else:
return_with_leftover_noise = "disable"
latent = nodes.KSamplerAdvanced().sample(model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise)[0]
latent = separated_sample(model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise)
return basic_pipe, latent, vae