| import random | |
| import comfy | |
| import comfy.model_management | |
| import comfy.sample | |
| import comfy.samplers | |
| import comfy.utils | |
| from comfy.utils import ProgressBar | |
| import numpy as np | |
| import torch | |
| from ..utils import VyroParams | |
| # from ..utils.restart_sampling import ( | |
| # DDIMWrapper, | |
| # KSamplerRestartWrapper, | |
| # RESWrapper, | |
| # SCHEDULER_MAPPING, | |
| # _restart_scheduler, | |
| # _restart_segments, | |
| # _total_steps, | |
| # prepare_restart_segments, | |
| # ) | |
| from ..utils.sdxl_ksampler import CfgMethods#, sdxl_ksampler | |
| def get_supported_samplers(): | |
| samplers = comfy.samplers.KSampler.SAMPLERS.copy() | |
| samplers.remove("uni_pc") | |
| samplers.remove("uni_pc_bh2") | |
| # SDE samplers cannot be used with restarts | |
| samplers.remove("dpmpp_sde") | |
| samplers.remove("dpmpp_sde_gpu") | |
| samplers.remove("dpmpp_2m_sde") | |
| samplers.remove("dpmpp_2m_sde_gpu") | |
| samplers.append("res") | |
| return samplers | |
| # def get_supported_restart_schedulers(): | |
| # return list(SCHEDULER_MAPPING.keys()) | |
| # def sdxl_restarts_ksampler(base_model, refiner_model, seed, base_steps, refiner_steps, cfg, sampler_name, scheduler, base_positive, base_negative, refiner_positive, refiner_negative, latent, segments, restart_scheduler, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, refiner_detail_boost=0.0,cfg_clamp_after_step=0): | |
| # global _total_steps, _restart_segments, _restart_scheduler | |
| # _restart_scheduler = restart_scheduler | |
| # _restart_segments = prepare_restart_segments(segments) | |
| # if sampler_name == "res": | |
| # sampler_wrapper = RESWrapper(start_step, last_step, _restart_segments, _restart_scheduler, cfg_clamp_after_step) | |
| # elif sampler_name == "ddim": | |
| # sampler_wrapper = DDIMWrapper(start_step, last_step, _restart_segments, _restart_scheduler, cfg_clamp_after_step) | |
| # else: | |
| # sampler_wrapper = KSamplerRestartWrapper(sampler_name, start_step, last_step, _restart_segments, _restart_scheduler, cfg_clamp_after_step) | |
| # # Add the additional steps to the progress bar | |
| # pbar_update_absolute = ProgressBar.update_absolute | |
| # def pbar_update_absolute_wrapper(self, value, total=None, preview=None): | |
| # pbar_update_absolute(self, value, _total_steps, preview) | |
| # ProgressBar.update_absolute = pbar_update_absolute_wrapper | |
| # try: | |
| # result = sdxl_ksampler(base_model, refiner_model, seed, base_steps, refiner_steps, cfg, sampler_name, | |
| # scheduler, base_positive, base_negative, refiner_positive, refiner_negative, | |
| # latent, denoise=denoise, disable_noise=False, | |
| # start_step=0, last_step=last_step, force_full_denoise=force_full_denoise, | |
| # dynamic_base_cfg=cfg, dynamic_refiner_cfg=cfg, | |
| # cfg_method=CfgMethods.TONEMAP, refiner_detail_boost=refiner_detail_boost, restart_wrapper=sampler_wrapper) | |
| # return result | |
| # finally: | |
| # sampler_wrapper.cleanup() | |
| # ProgressBar.update_absolute = pbar_update_absolute | |
| class VyroSDXLSampler: | |
| def INPUT_TYPES(s): | |
| return { | |
| "required": { | |
| "params": ("VYRO_PARAMS",), | |
| "base_model": ("MODEL",), | |
| "base_positive": ("CONDITIONING", ), | |
| "base_negative": ("CONDITIONING", ), | |
| "sampler_name": (get_supported_samplers(), ), | |
| "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), | |
| "base_ratio": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}), | |
| "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), | |
| "refiner_detail_boost": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05},), | |
| }, | |
| "optional": { | |
| "refiner_model": ("MODEL",), | |
| "refiner_positive": ("CONDITIONING", ), | |
| "refiner_negative": ("CONDITIONING", ), | |
| "latent_image": ("LATENT",), | |
| } | |
| } | |
| RETURN_TYPES = ("LATENT",) | |
| FUNCTION = "sample" | |
| CATEGORY = "Vyro/Samplers" | |
| def sample(self, params:VyroParams, base_model, base_positive, base_negative, sampler_name, scheduler, base_ratio, denoise, refiner_detail_boost, refiner_model=None, refiner_positive=None, refiner_negative=None, latent_image=None): | |
| if latent_image is None: | |
| latent_image = params.latents | |
| noise_seed = params.seed | |
| steps = params.steps | |
| cfg = params.cfg | |
| cfg_method = CfgMethods.TONEMAP | |
| dynamic_base_cfg = 0.0 | |
| dynamic_refiner_cfg = 0.0 | |
| has_refiner_model = refiner_model is not None | |
| base_steps = int(steps * (base_ratio + 0.0001)) if has_refiner_model else steps | |
| refiner_steps = max(0, steps - base_steps) | |
| if denoise < 0.005: | |
| return (params.latents,) | |
| torch.manual_seed(params.seed) | |
| random.seed(params.seed) | |
| np.random.seed(params.seed) | |
| if refiner_steps == 0 or not has_refiner_model: | |
| result = sdxl_ksampler(base_model, None, noise_seed, base_steps, 0, cfg, sampler_name, | |
| scheduler, base_positive, base_negative, None, None, | |
| latent_image, denoise=denoise, disable_noise=False, start_step=0, last_step=steps, | |
| force_full_denoise=True, dynamic_base_cfg=dynamic_base_cfg, cfg_method=cfg_method) | |
| else: | |
| result = sdxl_ksampler(base_model, refiner_model, noise_seed, base_steps, refiner_steps, cfg, sampler_name, | |
| scheduler, base_positive, base_negative, refiner_positive, refiner_negative, | |
| latent_image, denoise=denoise, disable_noise=False, | |
| start_step=0, last_step=steps, force_full_denoise=True, | |
| dynamic_base_cfg=dynamic_base_cfg, dynamic_refiner_cfg=dynamic_refiner_cfg, | |
| cfg_method=cfg_method, refiner_detail_boost=refiner_detail_boost) | |
| return result | |
| # class VyroKRestartSampler: | |
| # def __init__(self) -> None: | |
| # if 'res' not in comfy.samplers.KSampler.SAMPLERS: | |
| # comfy.samplers.KSampler.SAMPLERS.append('res') | |
| # @classmethod | |
| # def INPUT_TYPES(s): | |
| # return { | |
| # "required": { | |
| # "params": ("VYRO_PARAMS",), | |
| # "base_model": ("MODEL",), | |
| # "base_positive": ("CONDITIONING", ), | |
| # "base_negative": ("CONDITIONING", ), | |
| # "refiner_model": ("MODEL",), | |
| # "refiner_positive": ("CONDITIONING", ), | |
| # "refiner_negative": ("CONDITIONING", ), | |
| # "sampler_name": (get_supported_samplers(), ), | |
| # "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), | |
| # "segments": ("STRING", {"default": "[3,2,0.06,0.30],[3,1,0.30,0.59]", "multiline": False}), | |
| # "restart_scheduler": (get_supported_restart_schedulers(), ), | |
| # "begin_at_step": ("INT", {"default": 1, "min": 0, "max": 10000}), | |
| # "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}), | |
| # "add_noise": (["enable", "disable"], ), | |
| # "return_with_leftover_noise": (["disable", "enable"], ), | |
| # "base_ratio": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}), | |
| # "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), | |
| # "refiner_detail_boost": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05},), | |
| # "cfg_clamp_after_step": ("INT", {"default": 0, "min": 0, "max": 10000}), | |
| # }, | |
| # "optional": { | |
| # "refiner_prep_steps": ("INT", {"default": 0, "min": 0, "max": 10}), | |
| # "latent_image": ("LATENT",), | |
| # }, | |
| # } | |
| # RETURN_TYPES = ("LATENT",) | |
| # FUNCTION = "sample" | |
| # CATEGORY = "Vyro/Samplers" | |
| # def sample(self, params:VyroParams, base_model, base_positive, base_negative, refiner_model, refiner_positive, refiner_negative, sampler_name, scheduler, segments, restart_scheduler, begin_at_step, end_at_step, add_noise, return_with_leftover_noise, base_ratio, denoise, refiner_detail_boost, refiner_prep_steps=0, cfg_clamp_after_step=0, latent_image=None): | |
| # force_full_denoise = True | |
| # if return_with_leftover_noise == "enable": | |
| # force_full_denoise = False | |
| # disable_noise = False | |
| # if add_noise == "disable": | |
| # disable_noise = True | |
| # steps = params.steps | |
| # base_steps = int(steps * (base_ratio + 0.0001)) | |
| # refiner_steps = max(0, steps - base_steps) | |
| # if latent_image is None: | |
| # input_latent = latent_image = params.latents | |
| # else: | |
| # input_latent = latent_image | |
| # if denoise < 0.01: | |
| # return (latent_image, ) | |
| # # if refiner_prep_steps is not None: | |
| # # if refiner_prep_steps >= base_steps: | |
| # # refiner_prep_steps = base_steps - 1 | |
| # # if refiner_prep_steps > 0: | |
| # # start_at_step = refiner_prep_steps | |
| # # precondition_result = nodes.common_ksampler(refiner_model, params.seed + 2, steps, params.cfg, sampler_name, scheduler, refiner_positive, refiner_negative, latent_image, denoise=denoise, disable_noise=False, start_step=steps - refiner_prep_steps, last_step=steps, force_full_denoise=False) | |
| # # input_latent = precondition_result[0] | |
| # torch.manual_seed(params.seed) | |
| # random.seed(params.seed) | |
| # np.random.seed(params.seed) | |
| # if base_steps >= steps: | |
| # return sdxl_restarts_ksampler(base_model, None, params.seed, base_steps, refiner_steps, params.cfg, sampler_name, scheduler, base_positive, base_negative, refiner_positive, refiner_negative, input_latent, segments, restart_scheduler=restart_scheduler, denoise=denoise, disable_noise=disable_noise, start_step=begin_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise, refiner_detail_boost=refiner_detail_boost, cfg_clamp_after_step=cfg_clamp_after_step) | |
| # # return restart_sampling(base_model, params.seed, steps, params.cfg, sampler_name, scheduler, base_positive, base_negative, input_latent, denoise, segments, restart_scheduler, begin_at_step, end_at_step, disable_noise, force_full_denoise) | |
| # return sdxl_restarts_ksampler(base_model, refiner_model, params.seed, base_steps, refiner_steps, params.cfg, sampler_name, scheduler, base_positive, base_negative, refiner_positive, refiner_negative, input_latent, segments, restart_scheduler=restart_scheduler, denoise=denoise, disable_noise=disable_noise, start_step=begin_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise, refiner_detail_boost=refiner_detail_boost, cfg_clamp_after_step=cfg_clamp_after_step) | |
| NODE_CLASS_MAPPINGS = { | |
| # "Vyro KRestart Sampler": VyroKRestartSampler, | |
| "Vyro SDXL Sampler": VyroSDXLSampler, | |
| } | |
| NODE_DISPLAY_NAME_MAPPINGS = { | |
| # "VyroKRestartSampler": "Vyro KRestart Sampler", | |
| "VyroSDXLSampler": "Vyro SDXL Sampler", | |
| } | |