diff --git a/__init__.py b/__init__.py index 0093dee..8a21d2b 100644 --- a/__init__.py +++ b/__init__.py @@ -4,7 +4,6 @@ from . import clyb_Samplers from . import clyb_Schedulers clyb_Samplers.add_samplers() -clyb_Schedulers.add_schedulers() NODE_CLASS_MAPPINGS = { "ChromaNAG": chroma_NAG.ChromaNAG, diff --git a/clyb_Guidance.py b/clyb_Guidance.py index c31b152..dd11cc9 100644 --- a/clyb_Guidance.py +++ b/clyb_Guidance.py @@ -1,4 +1,5 @@ import torch +import torch.nn.functional as F import math import comfy import re @@ -26,16 +27,20 @@ class ClybGuidance: "var_rescale": ("BOOLEAN", {"default": False, "tooltip":"Whether we use torch.var (true) or torch.std (false) for rescaling."}), "scale_up_ratio": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip":"Weight of: Initiating CFG at guidance scale 1, increasing to your guidance scale in the middle of diffusion, and lower back to 1."}), "scale_up_shift": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.01, "tooltip":"Whether to shift to your CFG scale later (lower than 1.0) or earlier (higher than 1.0) in the schedule."}), + "atan2sin_ratio": ("FLOAT", {"default": 0.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip":"Applies standard deviation renormalization of CFG to cond at this rate."}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "sampling/custom_sampling" - def patch(self, model, eta, norm_threshold, momentum, momentum_beta, momentum_renorm, scalar_projection, scalar_logsumexp, rescale_phi, var_rescale, scale_up_ratio, scale_up_shift): + def patch(self, model, eta, norm_threshold, momentum, momentum_beta, momentum_renorm, scalar_projection, scalar_logsumexp, rescale_phi, var_rescale, scale_up_ratio, scale_up_shift, atan2sin_ratio): running_avg = 0 prev_sigma = None + m = model.clone() + model_sampling = m.get_model_object("model_sampling") + def pre_cfg_function(args): nonlocal running_avg, prev_sigma @@ -45,7 +50,6 @@ class ClybGuidance: uncond = args["conds_out"][1] sigma = args["sigma"][0] cond_scale = args["cond_scale"] - model_sampling = model.get_model_object("model_sampling") flow = False if isinstance(model_sampling, comfy.model_sampling.CONST): flow = True @@ -65,20 +69,22 @@ class ClybGuidance: else: timestep_ratio = model_sampling.timestep(args["timestep"]).float() / float(num_timesteps - 1) # Ratio scaling from 0 to 1 as diffusion goes on. - guidance_multiplier = torch.lerp(torch.ones_like(args["timestep"]), torch.sin((1. - timestep_ratio**scale_up_shift) * math.pi), weight=scale_up_ratio) # Lerp from static 1.0 scale to bell-curve (sine wave) scale - cfg_scalar = 1 / cond_scale + guidance_multiplier * ((cond_scale - 1) / cond_scale) # The guidance scale ought to be at least 1 (guidance is multiplied by cond scale, so ensure a min of 1/cond_scale) + if atan2sin_ratio != 0: + uncond = uncond.lerp(cond.atan().sin_().div_(uncond.atan().cos_()), weight=atan2sin_ratio) if scalar_projection: cond_flat, uncond_flat = cond.view(cond.shape[0], -1), uncond.view(uncond.shape[0], -1) - dot_product = torch.logsumexp(cond * uncond, dim=1, keepdim=True) if scalar_logsumexp else torch.sum(cond_flat * uncond_flat, dim=1, keepdim=True) + dot_product = torch.logsumexp(cond_flat * uncond_flat, dim=1, keepdim=True) if scalar_logsumexp else torch.sum(cond_flat * uncond_flat, dim=1, keepdim=True) - squared_norm = torch.logsumexp(uncond**2, dim=1, keepdim=True) if scalar_logsumexp else torch.sum(uncond_flat**2, dim=1, keepdim=True) + squared_norm = torch.logsumexp(uncond_flat**2, dim=1, keepdim=True) if scalar_logsumexp else torch.sum(uncond_flat**2, dim=1, keepdim=True) alpha = dot_product / squared_norm.clamp_min(1e-7) uncond = uncond * alpha + guidance_multiplier = torch.lerp(torch.ones_like(args["timestep"]), torch.sin((1. - timestep_ratio**scale_up_shift) * math.pi), weight=scale_up_ratio) # Lerp from static 1.0 scale to bell-curve (sine wave) scale + cfg_scalar = 1 / cond_scale + guidance_multiplier * ((cond_scale - 1) / cond_scale) # The guidance scale ought to be at least 1 (guidance is multiplied by cond scale, so ensure a min of 1/cond_scale) guidance = ((cond - uncond) * cfg_scalar) if scale_up_ratio != 0 else (cond - uncond) # Guidance is equivalent to (uncond -> cond) if momentum != 0: @@ -86,22 +92,19 @@ class ClybGuidance: running_avg = guidance else: running_avg = running_avg.lerp(guidance, weight=1. - momentum_beta)#running_avg.lerp(guidance, weight=1. - abs(momentum))# Update running average - #running_avg = running_avg * (guidance.pow(2).mean().sqrt_() / running_avg.pow(2).mean().sqrt_().clamp_min_(1e-8)) # Normalize running average to guidance momentumized_guidance = guidance.add(running_avg, alpha=momentum) momentumized_guidance_flat, guidance_flat = momentumized_guidance.view(momentumized_guidance.shape[0], -1), guidance.view(guidance.shape[0], -1) - guidance = momentumized_guidance.lerp(momentumized_guidance * (guidance_flat.norm(1, dim=1, keepdim=True) / momentumized_guidance_flat.norm(1, dim=1, keepdim=True).clamp_min(1e-7)), weight=momentum_renorm) - - if norm_threshold > 0: - guidance_norm = guidance.view(guidance.shape[0], -1).norm(p=2, dim=1, keepdim=True) - scale = torch.minimum( - torch.ones_like(guidance_norm), - norm_threshold / guidance_norm - ) - guidance = guidance * scale + guidance = momentumized_guidance.lerp(momentumized_guidance * (guidance_flat.norm(2, dim=1, keepdim=True) / momentumized_guidance_flat.norm(2, dim=1, keepdim=True).clamp_min(1e-7)), weight=momentum_renorm) guidance_parallel, guidance_orthogonal = project(guidance, cond) modified_guidance = guidance_orthogonal + eta * guidance_parallel + if norm_threshold > 0: + cond_norm = cond.norm(p=2, dim=tuple(range(1, len(cond.shape))), keepdim=True) * norm_threshold + guidance_norm = (uncond + modified_guidance * cond_scale).norm(p=2, dim=tuple(range(1, len(cond.shape))), keepdim=True) + if guidance_norm >= cond_norm: + modified_guidance = modified_guidance * (cond_norm / guidance_norm) + modified_cond = (uncond + modified_guidance) if rescale_phi != 0: # Formulate CFG @@ -122,146 +125,202 @@ class ClybGuidance: return [modified_cond, uncond] + args["conds_out"][2:] - m = model.clone() + """ + TODO: Rework these and add a selector + def magnitude_guidance(args): + cond = args["cond_denoised"] + uncond = args["uncond_denoised"] + cond_scale = args["cond_scale"] + x = args['input'] + out = args["denoised"] + + # Flatten cond and uncond, ensure cond scale is positive, utilize double precision + device = cond.device + b, c, h, w = cond.shape + + # 1. Create the 2D Hann window kernel for convolution + hann_1d = torch.signal.windows.hann(63, device=device) + #hann_2d = torch.outer(hann_1d, hann_1d) + # Normalize the kernel so that the sum of its elements is 1 + hann_1d /= hann_1d.sum() + + # Reshape kernel for depthwise convolution: (out_channels, in_channels/groups, kH, kW) + # We use groups=c to apply the same 2D filter to each channel independently. + kernel = hann_1d.unsqueeze(0).unsqueeze(0)#.repeat(cond.shape[1], 1, 1, 1) + + # 2. Calculate the local average magnitude of the `cond` tensor + # We use the absolute value to measure magnitude, not the raw value. + # 'same' padding ensures the output has the same HxW dimensions as the input. + view_shape = (cond.shape[0], -1) + cond_flat = cond.view(view_shape) + uncond_flat = uncond.view(view_shape) + local_avg_magnitude = F.conv1d((cond_flat - uncond_flat), kernel, padding='same') + + # 3. Normalize the magnitude map for each image in the batch to the [0, 1] range + # This makes the `strength` parameter behave consistently across different images. + batch_mins = torch.min(local_avg_magnitude.view(view_shape), dim=-1)[0]#.view(cond.shape[0], 1, 1, 1) + batch_maxs = torch.max(local_avg_magnitude.view(view_shape), dim=-1)[0]#.view(cond.shape[0], 1, 1, 1) + + normalized_magnitude = (local_avg_magnitude - batch_mins) / (batch_maxs - batch_mins).clamp_min_(1e-16) + + # 4. Create the local scale multiplier + # The dampening is proportional to the normalized local magnitude and the `strength` param. + # We clamp to ensure the multiplier stays within a reasonable [0, 1] range. + dampening = torch.clamp(normalized_magnitude, 0, 1.0) + + # Invert the dampening: high magnitude -> low multiplier, low magnitude -> high multiplier + scale_multiplier = 1.0 - dampening + + cond_normed = cond_flat / torch.linalg.norm(cond_flat,dim=1,keepdim=True) + uncond_normed = uncond_flat / torch.linalg.norm(uncond_flat,dim=1,keepdim=True) + dot_product = torch.sum(cond_normed*uncond_normed,dim=1,keepdim=True) + # The final local scale is the base scale modulated by our multiplier + local_scale = cond_scale * scale_multiplier * dot_product + + #print(local_scale) + + # 5. Apply the standard CFG formula, but with our dynamic local scale + # uncond + local_scale * (cond - uncond) + guided_tensor = cond + (local_scale * (cond_flat - uncond_flat)).view(cond.shape) + + output = out.lerp(guided_tensor.view(cond.shape).to(cond.dtype), weight=atan2sin_ratio) + + if norm_threshold > 0: + guidance_norm = output.norm(p=2, dim=1, keepdim=True) + output = torch.where( + guidance_norm > norm_threshold, + output * (norm_threshold / guidance_norm), + output + ) + return output + + def frequency_guidance(args): + cond = args["cond_denoised"] + uncond = args["uncond_denoised"] + cond_scale = args["cond_scale"] + x = args['input'] + out = args["denoised"] + + # 1. Move to Frequency domain using 2D Fast Fourier Transform + # We use norm='ortho' to ensure the transform is unitary and preserves energy. + fft_cond = torch.fft.fftshift(torch.fft.fftn(cond.to(torch.float64), norm='ortho')) + fft_uncond = torch.fft.fftshift(torch.fft.fftn(uncond.to(torch.float64), norm='ortho')) + fft_out = torch.fft.fftshift(torch.fft.fftn(out.to(torch.float64), norm='ortho')) + + # 1. Create the 2D Hann window kernel for convolution + hann_1d = torch.signal.windows.hann(5, device=cond.device) + #hann_2d = torch.outer(hann_1d, hann_1d) + # Normalize the kernel so that the sum of its elements is 1 + hann_1d /= hann_1d.sum() + + kernel = hann_1d.unsqueeze(0).unsqueeze(0)#.repeat(cond.shape[1], 1, 1, 1) + + # 2. Calculate the local average magnitude of the `cond` tensor + # We use the absolute value to measure magnitude, not the raw value. + # 'same' padding ensures the output has the same HxW dimensions as the input. + view_shape = (fft_cond.shape[0], -1) + + fft_cond_flat = fft_cond.view(view_shape) + fft_uncond_flat = fft_uncond.view(view_shape) + + fft_cond_real = fft_cond_flat.real + fft_uncond_real = fft_uncond_flat.real + #guidance_direction = (cond - uncond) + local_avg_magnitude = F.conv1d(fft_cond_real, kernel.to(torch.float64), padding='same') + + # 3. Normalize the magnitude map for each image in the batch to the [0, 1] range + # This makes the `strength` parameter behave consistently across different images. + #view_shape = (cond.shape[0], -1) + batch_mins = torch.min(local_avg_magnitude.view(view_shape), dim=-1)[0]#.view(cond.shape[0], 1, 1, 1) + batch_maxs = torch.max(local_avg_magnitude.view(view_shape), dim=-1)[0]#.view(cond.shape[0], 1, 1, 1) + + normalized_magnitude = (local_avg_magnitude - batch_mins) / (batch_maxs - batch_mins).clamp_min_(1e-16) + + dampening = torch.clamp(normalized_magnitude, 0, 1.0) + + # Invert the dampening: high magnitude -> low multiplier, low magnitude -> high multiplier + scale_multiplier = 1.0 - dampening + + # The final local scale is the base scale modulated by our multiplier and dot product + cond_normed = fft_cond_real / torch.linalg.norm(fft_cond_real,dim=1,keepdim=True).clamp_min_(1e-16) + uncond_normed = fft_uncond_real / torch.linalg.norm(fft_uncond_real,dim=1,keepdim=True).clamp_min_(1e-16) + dot_product = torch.sum(cond_normed*uncond_normed,dim=1,keepdim=True) + + local_scale = cond_scale * scale_multiplier * dot_product + + print(local_scale) + + guided_tensor = fft_cond + (local_scale.to(torch.cdouble) * (fft_cond_flat - fft_uncond_flat)).view(fft_cond.shape) + + guided_tensor = torch.fft.ifftshift(guided_tensor) + guided_tensor = torch.fft.ifftn(guided_tensor, norm='ortho').real + + output = out.lerp(guided_tensor.view(cond.shape).to(cond.dtype), weight=atan2sin_ratio) + + if norm_threshold > 0: + guidance_norm = output.norm(p=2, dim=1, keepdim=True) + output = torch.where( + guidance_norm > norm_threshold, + output * (norm_threshold / guidance_norm), + output + ) + return output + + def kron_guidance(args): + cond = args["cond_denoised"] + uncond = args["uncond_denoised"] + cond_scale = args["cond_scale"] + x = args['input'] + out = args["denoised"] + + guidance = cond - uncond + + device = guidance.device + + # 2. To Frequency Domain + # Apply 2D FFT to the spatial dimensions (H, W) + guidance_fft = torch.fft.fftn(guidance) + + # Shift the zero-frequency component to the center for easier mask creation + guidance_fft_shifted = torch.fft.fftshift(guidance_fft).reshape(guidance_fft.shape[0], -1) + + U, S, Vh = torch.linalg.svd(guidance_fft_shifted, full_matrices=False) + + # 3. Truncate the SVD components to the desired rank + U_k = U[:, :1] + S_k = S[:1] + Vh_k = Vh[:1, :] + #S_k_sqrt = S_k.sqrt() + + # 4. Apply the Mask + # Multiply the shifted FFT of the guidance by the scale mask + scaled_guidance_fft_shifted = (U_k @ (torch.diag(S_k).to(torch.cfloat) @ Vh_k)).reshape(guidance_fft.shape) + + # 5. Back to Latent Domain + # Inverse shift to move the zero-frequency back to the corner + scaled_guidance_fft = torch.fft.ifftshift(scaled_guidance_fft_shifted) + + # Inverse 2D FFT to get back to the spatial (latent) domain + # The result of ifft2 will be complex; we take the real part. The imaginary + # part should be negligible for real inputs. + modified_guidance = torch.fft.ifftn(scaled_guidance_fft).real + + # 6. Apply Guidance + # Add the modified guidance to the unconditional prediction + output = out + (guidance - modified_guidance) * atan2sin_ratio + + #output = out.lerp(guided_tensor.view(cond.shape).to(cond.dtype), weight=atan2sin_ratio) + + if norm_threshold > 0: + guidance_norm = output.norm(p=2, dim=1, keepdim=True) + output = torch.where( + guidance_norm > norm_threshold, + output * (norm_threshold / guidance_norm), + output + ) + return output + """ + m.set_model_sampler_pre_cfg_function(pre_cfg_function) - return (m,) -#args = {"conds":conds, "conds_out": out, "cond_scale": self.cfg, "timestep": timestep, -# "input": x, "sigma": timestep, "model": self.inner_model, "model_options": model_options} -#out = fn(args) - -def create_number_range(range_str: str) -> list[int] | None: - """ - Creates a list of numbers from a string in "start-end" format. - - Args: - range_str: The input string (e.g., "1-5", "10-20", " 5 - 10 "). - - Returns: - A list of integers representing the range (inclusive), - or None if the string format is invalid or start > end. - """ - # The core 're' module line to define and apply the pattern: - # 1. r"..." denotes a raw string to avoid issues with backslashes. - # 2. (\d+) is a capturing group for one or more digits (the start number). - # 3. \s* matches zero or more whitespace characters (optional spaces around the hyphen). - # 4. - matches the literal hyphen. - # 5. \s* matches zero or more whitespace characters again. - # 6. (\d+) is another capturing group for the end number. - # re.match() attempts to match the pattern from the beginning of the string. - match = re.match(r"(\d+)\s*-\s*(\d+)", range_str.strip()) - - if match: - # Extract the captured groups and convert them to integers - start_str, end_str = match.groups() - start = int(start_str) - end = int(end_str) - - # Ensure the start is not greater than the end for a valid range - if start <= end: - return list(range(start, end + 1)) - else: - # Handle cases like "5-1" if they should not produce a range - print(f"Warning: Start ({start}) is greater than end ({end}) for '{range_str}'") - return None # Or [] if an empty list is preferred for invalid ranges - else: - return None # String does not match the expected format - -class ClybLayerGuidanceDiT: - ''' - Enhance guidance towards detailed dtructure by having another set of CFG negative with skipped layers. - Inspired by Perturbed Attention Guidance (https://arxiv.org/abs/2403.17377) - Original experimental implementation for SD3 by Dango233@StabilityAI. - ''' - @classmethod - def INPUT_TYPES(s): - return {"required": {"model": ("MODEL", ), - "double_layers": ("STRING", {"default": "7, 8, 9", "multiline": False}), - "single_layers": ("STRING", {"default": "7, 8, 9", "multiline": False}), - "scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 10.0, "step": 0.1}), - "start_percent": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 1.0, "step": 0.001}), - "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}), - "rescaling_scale": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.01}), - "attn_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "skip_guidance" - EXPERIMENTAL = True - - DESCRIPTION = "Generic version of ClybLayerGuidance node that can be used on every DiT model." - - CATEGORY = "advanced/guidance" - - def skip_guidance(self, model, scale, start_percent, end_percent, double_layers="", single_layers="", rescaling_scale=0, attn_scale=1.0): - # check if layer is comma separated integers - def skip(args, extra_args): - print(f"ARGS: {args.items()}", "\n\n\n", f"EXTRA_ARGS: {extra_args.items()}", "\n\n\n") - for x in args: - if 'vec' in x: - for y in x: - if 'scale' in y: - args[x][y] = args[x][y] * attn_scale - #args[x] = y * attn_scale - #print(x, y) - return args - #for x, y in args.items(): - # if 'img' in x: - # return x - - model_sampling = model.get_model_object("model_sampling") - sigma_start = model_sampling.percent_to_sigma(start_percent) - sigma_end = model_sampling.percent_to_sigma(end_percent) - - #double_layers = re.findall(r'\d+', double_layers) - #double_layers = [int(i) for i in double_layers] - double_layers = create_number_range(double_layers) - - #single_layers = re.findall(r'\d+', single_layers) - #single_layers = [int(i) for i in single_layers] - single_layers = create_number_range(single_layers) - - if len(double_layers) == 0 and len(single_layers) == 0: - return (model, ) - - def post_cfg_function(args): - model = args["model"] - cond_pred = args["cond_denoised"] - cond = args["cond"] - cfg_result = args["denoised"] - sigma = args["sigma"] - x = args["input"] - model_options = args["model_options"].copy() - #print(model_options) - for layer in double_layers: - model_options = comfy.model_patcher.set_model_options_patch_replace(model_options, skip, "dit", "double_block", layer) - - for layer in single_layers: - model_options = comfy.model_patcher.set_model_options_patch_replace(model_options, skip, "dit", "single_block", layer) - - model_sampling.percent_to_sigma(start_percent) - - sigma_ = sigma[0].item() - if scale > 0 and sigma_ >= sigma_end and sigma_ <= sigma_start: - (slg,) = comfy.samplers.calc_cond_batch(model, [cond], x, sigma, model_options) - cfg_result = cfg_result + (cond_pred - slg) * scale - if rescaling_scale != 0: - factor = cond_pred.std() / cfg_result.std() - factor = rescaling_scale * factor + (1 - rescaling_scale) - cfg_result *= factor - - return cfg_result - - m = model.clone() - m.set_model_sampler_post_cfg_function(post_cfg_function) - - return (m, ) - -#NODE_CLASS_MAPPINGS = { -# "ClybGuidance": ClybGuidance, -# "ClybLayerGuidanceDiT": ClybLayerGuidanceDiT, -#} - -#NODE_DISPLAY_NAME_MAPPINGS = { -# "ClybGuidance": "ClybGuidance", -# "ClybLayerGuidanceDiT": "ClybLayerGuidanceDiT", -#} + return (m,) \ No newline at end of file diff --git a/clyb_Samplers.py b/clyb_Samplers.py index 6674452..d794105 100644 --- a/clyb_Samplers.py +++ b/clyb_Samplers.py @@ -7,7 +7,7 @@ from comfy.k_diffusion.sampling import default_noise_sampler import comfy.samplers @torch.no_grad() -def sampler_clyb_bdf(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, flow=False): +def sampler_clyb_bdf(model, x, sigmas, extra_args=None, callback=None, disable=None, scalar="atan2sin+projection", eta=1., s_noise=1., noise_sampler=None, flow=False): extra_args = {} if extra_args is None else extra_args seed = extra_args.get("seed", None) noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler @@ -34,14 +34,25 @@ def sampler_clyb_bdf(model, x, sigmas, extra_args=None, callback=None, disable=N alpha_down = 1 - sigma_down renoise_coeff = (sigmas[i+1]**2 - sigma_down**2*alpha_ip1**2/alpha_down**2)**0.5 - first_denoised = prev_denoised if prev_denoised is not None else model(x, sigmas[i] * s_in, **extra_args) + first_denoised = prev_denoised if (prev_denoised is not None and sigma_down > 0) else model(x, sigmas[i] * s_in, **extra_args) - if sigma_down > 0: + if sigma_down > 0 and i > 0: x_faux = first_denoised.lerp(x, weight=sigma_down/sigmas[i]) denoised2 = model(x_faux, sigma_down * s_in, **extra_args) second_denoised = (first_denoised + denoised2) / 2 - scalar = (denoised2 * second_denoised) / (second_denoised.pow(2).clamp_min(1e-6)) - denoised_prime = second_denoised * scalar + match scalar: + case "projection": + scaling = (denoised2 * second_denoised) / (second_denoised.pow(2).clamp_min(1e-7)) + denoised_prime = second_denoised * scaling + case "atan2sin": + denoised_prime = denoised2.atan().sin_().div_(second_denoised.atan().cos_()) + case "atan2sin+projection": + denoised_prime = denoised2.atan().sin_().div_(second_denoised.atan().cos_()) + scaling = (denoised2 * denoised_prime) / (denoised_prime.pow(2).clamp_min(1e-7)) + denoised_prime = denoised_prime * scaling + case _: + scaling = (denoised2 * second_denoised) / (second_denoised.pow(2).clamp_min(1e-7)) + denoised_prime = second_denoised * scaling else: denoised_prime = first_denoised @@ -61,11 +72,11 @@ def sampler_clyb_bdf(model, x, sigmas, extra_args=None, callback=None, disable=N return x @torch.no_grad() -def sample_clyb_bdf(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): +def sample_clyb_bdf(model, x, sigmas, extra_args=None, callback=None, disable=None, scalar="atan2sin+projection", eta=1., s_noise=1., noise_sampler=None): flow = False if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST): flow = True - return sampler_clyb_bdf(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, flow=flow) + return sampler_clyb_bdf(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, scalar=scalar, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, flow=flow) # The following function adds the samplers during initialization, in __init__.py def add_samplers(): @@ -95,8 +106,10 @@ discard_penultimate_sigma_samplers = set(()) class SamplerClyb_BDF: @classmethod def INPUT_TYPES(s): + NOISE_SAMPLER_NAMES=("projection", "atan2sin", "atan2sin+projection") return {"required": - {"eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}), + {"scalar": (NOISE_SAMPLER_NAMES, {"default": NOISE_SAMPLER_NAMES[2]}), + "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}), "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}), } } @@ -105,6 +118,6 @@ class SamplerClyb_BDF: FUNCTION = "get_sampler" - def get_sampler(self, eta, s_noise): - sampler = comfy.samplers.ksampler("clyb_bdf", {"eta": eta, "s_noise": s_noise}) + def get_sampler(self, scalar, eta, s_noise): + sampler = comfy.samplers.ksampler("clyb_bdf", {"scalar": scalar, "eta": eta, "s_noise": s_noise}) return (sampler, ) \ No newline at end of file diff --git a/clyb_Schedulers.py b/clyb_Schedulers.py index 1b8ef18..7d35456 100644 --- a/clyb_Schedulers.py +++ b/clyb_Schedulers.py @@ -1,5 +1,6 @@ import numpy import torch +from comfy.samplers import SchedulerHandler, SCHEDULER_HANDLERS, SCHEDULER_NAMES def inverse_squared_scheduler(model_sampling, steps): total_timesteps = (len(model_sampling.sigmas) - 1) @@ -55,21 +56,12 @@ class PrintSigmas: print(sigmas) return (sigmas, ) -def add_schedulers(): - from comfy.samplers import KSampler, k_diffusion_sampling - added = 0 - for scheduler in extra_schedulers: #getattr(self, "sample_{}".format(extra_samplers)) - if scheduler not in KSampler.SCHEDULERS: - try: - idx = KSampler.SCHEDULERS.index("ddim_uniform") # Last item in the samplers list - KSampler.SCHEDULERS.insert(idx+1, scheduler) # Add our custom samplers - setattr(k_diffusion_sampling, "get_sigmas_{}".format(scheduler), extra_schedulers[scheduler]) - added += 1 - except ValueError as err: - pass - if added > 0: - import importlib - importlib.reload(k_diffusion_sampling) +scheduler_name = "inverse_squared" +if scheduler_name not in SCHEDULER_HANDLERS: + scheduler_handler = SchedulerHandler(handler=inverse_squared_scheduler, use_ms=True) + SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler + if scheduler_name not in SCHEDULER_NAMES: + SCHEDULER_NAMES.append(scheduler_name) extra_schedulers = { "inverse_squared": inverse_squared_scheduler,