diff --git a/__init__.py b/__init__.py index c78fb38..ab9f1b3 100644 --- a/__init__.py +++ b/__init__.py @@ -5,9 +5,10 @@ settings.load_settings() if settings.SETTINGS.btp_enabled: from .py import betterTaesdPreview # noqa: F401 -from .py import ( - deepshrink, - hypertile, +from .py.nodes import ( + blockOps, + deepShrink, + hyperTile, modelPatchConditional, refinerAfter, samplers, @@ -15,14 +16,15 @@ from .py import ( ) NODE_CLASS_MAPPINGS = { - "BlehHyperTile": hypertile.HyperTileBleh, - "BlehDeepShrink": deepshrink.DeepShrinkBleh, + "BlehBlockOps": blockOps.BlehBlockOps, + "BlehDeepShrink": deepShrink.DeepShrinkBleh, "BlehDiscardPenultimateSigma": sigmas.DiscardPenultimateSigma, - "BlehInsaneChainSampler": samplers.BlehInsaneChainSampler, "BlehForceSeedSampler": samplers.BlehForceSeedSampler, + "BlehHyperTile": hyperTile.HyperTileBleh, + "BlehInsaneChainSampler": samplers.BlehInsaneChainSampler, + "BlehLatentUpscaleBy": blockOps.BlehLatentUpscaleBy, "BlehModelPatchConditional": modelPatchConditional.ModelPatchConditionalNode, "BlehRefinerAfter": refinerAfter.BlehRefinerAfter, - "BlehLatentUpscaleBy": deepshrink.BlehLatentUpscaleBy, } NODE_DISPLAY_NAME_MAPPINGS = { diff --git a/py/deepshrink.py b/py/deepshrink.py deleted file mode 100644 index dc7e655..0000000 --- a/py/deepshrink.py +++ /dev/null @@ -1,403 +0,0 @@ -# Adapted from the ComfyUI built-in node - -import bisect -import importlib as il - -import comfy -import numpy as np -import torch - -dare = il.import_module("custom_nodes.ComfyUI-DareMerge.ddare.merge") - -UPSCALE_METHODS = ( - "bicubic", - "nearest-exact", - "bilinear", - "area", - "bislerp", - "bislerp_alt", - "colorize", - "hslerp", - "bibislerp", -) - -FILTER_SIZES = ( - np.array([1.0]), - np.array([1.0, 1.0]), - np.array([1.0, 2.0, 1.0]), - np.array([1.0, 3.0, 3.0, 1.0]), - np.array([1.0, 4.0, 6.0, 4.0, 1.0]), - np.array([1.0, 5.0, 10.0, 10.0, 5.0, 1.0]), - np.array([1.0, 6.0, 15.0, 20.0, 15.0, 6.0, 1.0]), -) - - -def make_filter(channels, dtype, size=3): - a = FILTER_SIZES[size - 1] - filt = torch.tensor(a[:, None] * a[None, :], dtype=dtype) - filt = filt / torch.sum(filt) - return filt[None, None, :, :].repeat((channels, 1, 1, 1)) - - -def scale_samples( - samples, - width, - height, - mode="bicubic", - mode_h=None, - antialias_size=0, -): - if mode_h is None: - mode_h = mode - if mode in ("bicubic", "nearest-exact", "bilinear", "area"): - result = torch.nn.functional.interpolate( - samples, - size=(height, width), - mode=mode, - antialias=antialias_size > 7, - ) - - result = biderp(samples, width, height, mode, mode_h) - if antialias_size < 0 or antialias_size > 7: - return result - channels = result.shape[1] - filt = make_filter(channels, result.dtype, antialias_size).to(result.device) - return torch.nn.functional.conv2d(result, filt, groups=channels, padding="same") - - -# Modified from ComfyUI -def biderp(samples, width, height, mode="bislerp", mode_h=None): - def slerp_orig(b1, b2, r): - c = b1.shape[-1] - - # norms - b1_norms = torch.norm(b1, dim=-1, keepdim=True) - b2_norms = torch.norm(b2, dim=-1, keepdim=True) - - # normalize - b1_normalized = b1 / b1_norms - b2_normalized = b2 / b2_norms - - # zero when norms are zero - b1_normalized[b1_norms.expand(-1, c) == 0.0] = 0.0 - b2_normalized[b2_norms.expand(-1, c) == 0.0] = 0.0 - - # slerp - dot = (b1_normalized * b2_normalized).sum(1) - omega = torch.acos(dot) - so = torch.sin(omega) - - # technically not mathematically correct, but more pleasing? - res = (torch.sin((1.0 - r.squeeze(1)) * omega) / so).unsqueeze( - 1, - ) * b1_normalized + (torch.sin(r.squeeze(1) * omega) / so).unsqueeze( - 1, - ) * b2_normalized - res *= (b1_norms * (1.0 - r) + b2_norms * r).expand(-1, c) - - # edge cases for same or polar opposites - res[dot > 1 - 1e-5] = b1[dot > 1 - 1e-5] - res[dot < 1e-5 - 1] = (b1 * (1.0 - r) + b2 * r)[dot < 1e-5 - 1] - return res - - # Modified from DARE merge - def hslerp(a, b, t): - interpolation_tensor = torch.zeros( - 1, - a.shape[-1], - device=a.device, - dtype=a.dtype, - ) - - interpolation_tensor[0, 0] = 1.0 - - result = (1 - t) * a + t * b - norm = (torch.norm(b - a, dim=1, keepdim=True) / 6) * interpolation_tensor - norm[t < 0.5, None] *= -1 - result += norm - return result - - if mode_h is None: - mode_h = mode - - modes = { - "bislerp_alt": dare.merge_tensors_slerp, - "colorize": dare.merge_tensors_colorize, - "hslerp": hslerp, - "bislerp": slerp_orig, - "bibislerp": dare.merge_tensors_bislerp, - } - derp_w, derp_h = modes.get(mode, slerp_orig), modes.get(mode_h, slerp_orig) - - def generate_bilinear_data(length_old, length_new, device): - coords_1 = torch.arange(length_old, dtype=torch.float32, device=device).reshape( - (1, 1, 1, -1), - ) - coords_1 = torch.nn.functional.interpolate( - coords_1, - size=(1, length_new), - mode="bilinear", - ) - ratios = coords_1 - coords_1.floor() - coords_1 = coords_1.to(torch.int64) - - coords_2 = ( - torch.arange(length_old, dtype=torch.float32, device=device).reshape( - (1, 1, 1, -1), - ) - + 1 - ) - coords_2[:, :, :, -1] -= 1 - coords_2 = torch.nn.functional.interpolate( - coords_2, - size=(1, length_new), - mode="bilinear", - ) - coords_2 = coords_2.to(torch.int64) - return ratios, coords_1, coords_2 - - orig_dtype = samples.dtype - samples = samples.float() - n, c, h, w = samples.shape - h_new, w_new = (height, width) - - # linear w - ratios, coords_1, coords_2 = generate_bilinear_data(w, w_new, samples.device) - coords_1 = coords_1.expand((n, c, h, -1)) - coords_2 = coords_2.expand((n, c, h, -1)) - ratios = ratios.expand((n, 1, h, -1)) - - pass_1 = samples.gather(-1, coords_1).movedim(1, -1).reshape((-1, c)) - pass_2 = samples.gather(-1, coords_2).movedim(1, -1).reshape((-1, c)) - ratios = ratios.movedim(1, -1).reshape((-1, 1)) - - result = derp_w(pass_1, pass_2, ratios) - result = result.reshape(n, h, w_new, c).movedim(-1, 1) - - # linear h - ratios, coords_1, coords_2 = generate_bilinear_data(h, h_new, samples.device) - coords_1 = coords_1.reshape((1, 1, -1, 1)).expand((n, c, -1, w_new)) - coords_2 = coords_2.reshape((1, 1, -1, 1)).expand((n, c, -1, w_new)) - ratios = ratios.reshape((1, 1, -1, 1)).expand((n, 1, -1, w_new)) - - pass_1 = result.gather(-2, coords_1).movedim(1, -1).reshape((-1, c)) - pass_2 = result.gather(-2, coords_2).movedim(1, -1).reshape((-1, c)) - ratios = ratios.movedim(1, -1).reshape((-1, 1)) - - result = derp_h(pass_1, pass_2, ratios) - result = result.reshape(n, h_new, w_new, c).movedim(-1, 1) - return result.to(orig_dtype) - - -class DeepShrinkBleh: - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - CATEGORY = "bleh/model_patches" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "model": ("MODEL",), - "commasep_block_numbers": ( - "STRING", - { - "default": "3", - }, - ), - "downscale_w": ( - "FLOAT", - {"default": 2.0, "min": 0.01, "max": 32.0, "step": 0.1}, - ), - "downscale_h": ( - "FLOAT", - {"default": 2.0, "min": 0.01, "max": 32.0, "step": 0.1}, - ), - "start_percent": ( - "FLOAT", - {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}, - ), - "start_fadeout_percent": ( - "FLOAT", - {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}, - ), - "end_percent": ( - "FLOAT", - {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.001}, - ), - "downscale_after_skip": ("BOOLEAN", {"default": True}), - "downscale_method_w": (UPSCALE_METHODS,), - "downscale_method_h": (UPSCALE_METHODS,), - "upscale_method_w": (UPSCALE_METHODS,), - "upscale_method_h": (UPSCALE_METHODS,), - "antialias_downscale_size": ("INT", {"default": 0, "max": 7}), - "antialias_upscale_size": ("INT", {"default": 0, "max": 7}), - }, - } - - def patch( - self, - model, - commasep_block_numbers, - downscale_w, - downscale_h, - start_percent, - start_fadeout_percent, - end_percent, - downscale_after_skip, - downscale_method_w, - downscale_method_h, - upscale_method_w, - upscale_method_h, - antialias_downscale_size, - antialias_upscale_size, - ): - # if downscale_factor == 1 or start_percent == 1 or end_percent == 0: - # return (model.clone(),) - block_numbers = tuple( - int(x) for x in commasep_block_numbers.split(",") if x.strip() - ) - # if downscale_factor < 1 and start_fadeout_percent != 1: - # raise ValueError( - # "BlehDeepShrink: Setting scale factor below 1 cannot be combined with fadeout", - # ) - downscale_w, downscale_h = 1.0 / downscale_w, 1.0 / downscale_h - if not (block_numbers and all(val > 0 and val <= 32 for val in block_numbers)): - raise ValueError( - "BlehDeepShrink: Bad value for block numbers: must be comma-separated list of numbers between 1-32", - ) - # antialias_downscale = antialias_downscale and downscale_method in ( - # "bicubic", - # "bilinear", - # ) - # antialias_upscale = antialias_upscale and upscale_method in ( - # "bicubic", - # "bilinear", - # ) - if start_fadeout_percent < start_percent: - start_fadeout_percent = start_percent - elif start_fadeout_percent > end_percent: - # No fadeout. - start_fadeout_percent = 1000.0 - - sigma_start = model.model.model_sampling.percent_to_sigma( - max(0.001, start_percent), - ) - sigma_end = model.model.model_sampling.percent_to_sigma(end_percent) - - # Arbitrary number that should have good enough precision - pct_steps = 400 - pct_incr = 1.0 / pct_steps - sig2pct = tuple( - model.model.model_sampling.percent_to_sigma(x / pct_steps) - for x in range(pct_steps, -1, -1) - ) - # print(sig2pct) - - def input_block_patch(h, transformer_options): - sigma = transformer_options["sigmas"][0].item() - if ( - sigma > sigma_start - or sigma < sigma_end - or transformer_options["block"][1] not in block_numbers - ): - return h - # This is obviously terrible but I couldn't find a better way to get the percentage from the current sigma. - idx = bisect.bisect_right(sig2pct, sigma) - if idx >= len(sig2pct): - # Sigma out of range somehow? - return h - pct = pct_incr * (pct_steps - idx) - if ( - pct < start_fadeout_percent - or start_fadeout_percent > end_percent - or pct > end_percent - ): - scaled_scale_w, scaled_scale_h = downscale_w, downscale_h - else: - # May or not be accurate but the idea is to scale the downscale factor by the percentage - # of the start fade to end deep shrink we have currently traversed. It at least sort of works. - downscale_pct = 1.0 - ( - (pct - start_fadeout_percent) - / (end_percent - start_fadeout_percent) - ) - scaled_scale_w = 1.0 - ((1.0 - downscale_w) * downscale_pct) - scaled_scale_h = 1.0 - ((1.0 - downscale_h) * downscale_pct) - orig_width, orig_height = h.shape[-1], h.shape[-2] - width, height = ( - round(orig_width * scaled_scale_w), - round(orig_height * scaled_scale_h), - ) - # if scaled_scale >= 0.98 or width >= orig_width or height >= orig_height: - # return h - if (0.98 <= scaled_scale_w <= 1.02 and 0.98 <= scaled_scale_h <= 1.02) or ( - width == orig_width and height == orig_height - ): - return h - return scale_samples( - h, - width, - height, - mode=downscale_method_w, - mode_h=downscale_method_h, - antialias_size=antialias_downscale_size, - ) - - def output_block_patch(h, hsp, _transformer_options): - if h.shape[2] == hsp.shape[2] and h.shape[3] == hsp.shape[3]: - return h, hsp - return scale_samples( - h, - hsp.shape[-1], - hsp.shape[-2], - mode=upscale_method_w, - mode_h=upscale_method_h, - antialias_size=antialias_upscale_size, - ), hsp - - m = model.clone() - if start_percent >= 1.0: - return (m,) - if downscale_after_skip: - m.set_model_input_block_patch_after_skip(input_block_patch) - else: - m.set_model_input_block_patch(input_block_patch) - m.set_model_output_block_patch(output_block_patch) - return (m,) - - -class BlehLatentUpscaleBy: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "samples": ("LATENT",), - "upscale_method": (UPSCALE_METHODS,), - "scale_width": ( - "FLOAT", - {"default": 1.5, "min": 0.01, "max": 8.0, "step": 0.01}, - ), - "scale_height": ( - "FLOAT", - {"default": 1.5, "min": 0.01, "max": 8.0, "step": 0.01}, - ), - "antialias": ("BOOLEAN", {"default": False}), - }, - } - - RETURN_TYPES = ("LATENT",) - FUNCTION = "upscale" - - CATEGORY = "latent" - - def upscale(self, samples, upscale_method, scale_width, scale_height): - s = samples.copy() - width = round(samples["samples"].shape[3] * scale_width) - height = round(samples["samples"].shape[2] * scale_height) - s["samples"] = scale_samples( - samples["samples"], - width, - height, - upscale_method, - ) - return (s,) diff --git a/py/nodes/blockOps.py b/py/nodes/blockOps.py new file mode 100644 index 0000000..5f12b02 --- /dev/null +++ b/py/nodes/blockOps.py @@ -0,0 +1,721 @@ +# Adapted from the ComfyUI built-in node + +import bisect +import importlib as il +import math +from enum import Enum, auto + +import numpy as np +import torch +import yaml +from torch import fft + +dare = il.import_module("custom_nodes.ComfyUI-DareMerge.ddare.merge") + + +def normalize(latent, target_min=None, target_max=None): + min_val = latent.min() + max_val = latent.max() + + if target_min is None: + target_min = min_val + if target_max is None: + target_max = max_val + + normalized = (latent - min_val) / (max_val - min_val) + return normalized * (target_max - target_min) + target_min + + +def hslerp(a, b, t): + if a.shape != b.shape: + raise ValueError("Input tensors a and b must have the same shape.") + + num_channels = a.size(1) + + interpolation_tensor = torch.zeros( + 1, + num_channels, + 1, + 1, + device=a.device, + dtype=a.dtype, + ) + interpolation_tensor[0, 0, 0, 0] = 1.0 + + result = (1 - t) * a + t * b + + if t < 0.5: + result += (torch.norm(b - a, dim=1, keepdim=True) / 6) * interpolation_tensor + else: + result -= (torch.norm(b - a, dim=1, keepdim=True) / 6) * interpolation_tensor + + return result + + +blending_modes = { + # Args: + # - a (tensor): Latent input 1 + # - b (tensor): Latent input 2 + # - t (float): Blending factor + # Interpolates between tensors a and b using normalized linear interpolation. + "bislerp": lambda a, b, t: normalize((1 - t) * a + t * b), + # Transfer the color from `b` to `a` by t` factor + "colorize": lambda a, b, t: a + (b - a) * t, + # Interpolates between tensors a and b using cosine interpolation. + "cosinterp": lambda a, b, t: ( + a + b - (a - b) * torch.cos(t * torch.tensor(math.pi)) + ) + / 2, + # Interpolates between tensors a and b using cubic interpolation. + "cuberp": lambda a, b, t: a + (b - a) * (3 * t**2 - 2 * t**3), + # Interpolates between tensors a and b using normalized linear interpolation, + # with a twist when t is greater than or equal to 0.5. + "hslerp": hslerp, + # Adds tensor b to tensor a, scaled by t. + "inject": lambda a, b, t: a + b * t, + # Interpolates between tensors a and b using linear interpolation. + "lerp": lambda a, b, t: (1 - t) * a + t * b, + # Simulates a brightening effect by adding tensor b to tensor a, scaled by t. + "lineardodge": lambda a, b, t: normalize(a + b * t), +} + +mscales = { + "none": (), + "bandpass": ( + (5, 0.0), # Low-pass filter + (15, 1.0), # Pass-through filter (allows mid-range frequencies) + (25, 0.0), # High-pass filter + ), + "lowpass": ( + ( + 10, + 1.0, + ), + ), # Allows low-frequency components, suppresses high-frequency components + "highpass": ( + ( + 10, + 0.0, + ), + ), # Suppresses low-frequency components, allows high-frequency components + "passthrough": ((10, 1.0),), # Passes all frequencies unchanged, no filtering + "gaussianblur": ( + ( + 10, + 0.5, + ), # Blurs the image by allowing a range of frequencies with a Gaussian shape + ), + "edge": ( + ( + 10, + 2.0, + ) + ), # Enhances edges and high-frequency features while suppressing low-frequency details + "sharpen": ( + ( + 10, + 1.5, + ) + ), # Increases the sharpness of the image by emphasizing high-frequency components + "multilowpass": ((5, 1.0), (10, 0.5), (15, 0.2)), # Multi-scale low-pass filter + "multihighpass": ((5, 0.0), (10, 0.5), (15, 0.8)), # Multi-scale high-pass filter + "multipassthrough": ( + (5, 1.0), + (10, 1.0), + (15, 1.0), + ), # Pass-through at different scales + "multigaussianblur": ((5, 0.5), (10, 0.8), (15, 0.2)), # Multi-scale Gaussian blur + "multiedge": ((5, 1.2), (10, 1.5), (15, 2.0)), # Multi-scale edge enhancement + "multisharpen": ((5, 1.5), (10, 2.0), (15, 2.5)), # Multi-scale sharpening +} + +UPSCALE_METHODS = ( + "bicubic", + "nearest-exact", + "bilinear", + "area", + "bislerp", + "bislerp_alt", + "colorize", + "hslerp", + "bibislerp", + "cosinterp", + "cuberp", + "inject", + "lerp", + "lineardodge", +) + +FILTER_SIZES = ( + np.array([1.0]), + np.array([1.0, 1.0]), + np.array([1.0, 2.0, 1.0]), + np.array([1.0, 3.0, 3.0, 1.0]), + np.array([1.0, 4.0, 6.0, 4.0, 1.0]), + np.array([1.0, 5.0, 10.0, 10.0, 5.0, 1.0]), + np.array([1.0, 6.0, 15.0, 20.0, 15.0, 6.0, 1.0]), +) + + +def make_filter(channels, dtype, size=3): + a = FILTER_SIZES[size - 1] + filt = torch.tensor(a[:, None] * a[None, :], dtype=dtype) + filt = filt / torch.sum(filt) + return filt[None, None, :, :].repeat((channels, 1, 1, 1)) + + +def scale_samples( + samples, + width, + height, + mode="bicubic", + mode_h=None, + antialias_size=0, +): + if mode_h is None: + mode_h = mode + if mode in ("bicubic", "nearest-exact", "bilinear", "area"): + result = torch.nn.functional.interpolate( + samples, + size=(height, width), + mode=mode, + antialias=antialias_size > 7, + ) + else: + result = biderp(samples, width, height, mode, mode_h) + if antialias_size < 1 or antialias_size > 7: + return result + channels = result.shape[1] + filt = make_filter(channels, result.dtype, antialias_size).to(result.device) + return torch.nn.functional.conv2d(result, filt, groups=channels, padding="same") + + +# Modified from ComfyUI +def biderp(samples, width, height, mode="bislerp", mode_h=None): + def slerp_orig(b1, b2, r): + c = b1.shape[-1] + + # norms + b1_norms = torch.norm(b1, dim=-1, keepdim=True) + b2_norms = torch.norm(b2, dim=-1, keepdim=True) + + # normalize + b1_normalized = b1 / b1_norms + b2_normalized = b2 / b2_norms + + # zero when norms are zero + b1_normalized[b1_norms.expand(-1, c) == 0.0] = 0.0 + b2_normalized[b2_norms.expand(-1, c) == 0.0] = 0.0 + + # slerp + dot = (b1_normalized * b2_normalized).sum(1) + omega = torch.acos(dot) + so = torch.sin(omega) + + # technically not mathematically correct, but more pleasing? + res = (torch.sin((1.0 - r.squeeze(1)) * omega) / so).unsqueeze( + 1, + ) * b1_normalized + (torch.sin(r.squeeze(1) * omega) / so).unsqueeze( + 1, + ) * b2_normalized + res *= (b1_norms * (1.0 - r) + b2_norms * r).expand(-1, c) + + # edge cases for same or polar opposites + res[dot > 1 - 1e-5] = b1[dot > 1 - 1e-5] + res[dot < 1e-5 - 1] = (b1 * (1.0 - r) + b2 * r)[dot < 1e-5 - 1] + return res + + # Modified from DARE merge + def hslerp(a, b, t): + interpolation_tensor = torch.zeros( + 1, + a.shape[-1], + device=a.device, + dtype=a.dtype, + ) + + interpolation_tensor[0, 0] = 1.0 + + result = (1 - t) * a + t * b + norm = (torch.norm(b - a, dim=1, keepdim=True) / 6) * interpolation_tensor + norm[t < 0.5, None] *= -1 + result += norm + return result + + if mode_h is None: + mode_h = mode + + modes = { + "bislerp_alt": dare.merge_tensors_slerp, + "colorize": blending_modes["colorize"], + "hslerp": hslerp, + "bislerp": slerp_orig, + "bibislerp": blending_modes["bislerp"], + "inject": blending_modes["inject"], + "lerp": blending_modes["lerp"], + "lineardodge": blending_modes["lineardodge"], + "cosinterp": blending_modes["cosinterp"], + "cuberp": blending_modes["cuberp"], + } + derp_w, derp_h = modes.get(mode, slerp_orig), modes.get(mode_h, slerp_orig) + + def generate_bilinear_data(length_old, length_new, device): + coords_1 = torch.arange(length_old, dtype=torch.float32, device=device).reshape( + (1, 1, 1, -1), + ) + coords_1 = torch.nn.functional.interpolate( + coords_1, + size=(1, length_new), + mode="bilinear", + ) + ratios = coords_1 - coords_1.floor() + coords_1 = coords_1.to(torch.int64) + + coords_2 = ( + torch.arange(length_old, dtype=torch.float32, device=device).reshape( + (1, 1, 1, -1), + ) + + 1 + ) + coords_2[:, :, :, -1] -= 1 + coords_2 = torch.nn.functional.interpolate( + coords_2, + size=(1, length_new), + mode="bilinear", + ) + coords_2 = coords_2.to(torch.int64) + return ratios, coords_1, coords_2 + + orig_dtype = samples.dtype + samples = samples.float() + n, c, h, w = samples.shape + h_new, w_new = (height, width) + + # linear w + ratios, coords_1, coords_2 = generate_bilinear_data(w, w_new, samples.device) + coords_1 = coords_1.expand((n, c, h, -1)) + coords_2 = coords_2.expand((n, c, h, -1)) + ratios = ratios.expand((n, 1, h, -1)) + + pass_1 = samples.gather(-1, coords_1).movedim(1, -1).reshape((-1, c)) + pass_2 = samples.gather(-1, coords_2).movedim(1, -1).reshape((-1, c)) + ratios = ratios.movedim(1, -1).reshape((-1, 1)) + + result = derp_w(pass_1, pass_2, ratios) + result = result.reshape(n, h, w_new, c).movedim(-1, 1) + + # linear h + ratios, coords_1, coords_2 = generate_bilinear_data(h, h_new, samples.device) + coords_1 = coords_1.reshape((1, 1, -1, 1)).expand((n, c, -1, w_new)) + coords_2 = coords_2.reshape((1, 1, -1, 1)).expand((n, c, -1, w_new)) + ratios = ratios.reshape((1, 1, -1, 1)).expand((n, 1, -1, w_new)) + + pass_1 = result.gather(-2, coords_1).movedim(1, -1).reshape((-1, c)) + pass_2 = result.gather(-2, coords_2).movedim(1, -1).reshape((-1, c)) + ratios = ratios.movedim(1, -1).reshape((-1, 1)) + + result = derp_h(pass_1, pass_2, ratios) + result = result.reshape(n, h_new, w_new, c).movedim(-1, 1) + return result.to(orig_dtype) + + +def ffilter(x, threshold, scale, scales=None, strength=1.0): + # FFT + if isinstance(x, list): + x = x[0] + if not isinstance(x, torch.Tensor): + return x + x_freq = fft.fftn(x.float(), dim=(-2, -1)) + x_freq = fft.fftshift(x_freq, dim=(-2, -1)) + + _batch, _channels, height, width = x_freq.shape + mask = torch.ones(x_freq.shape, device=x.device) + + crow, ccol = height // 2, width // 2 + mask[ + ..., + crow - threshold : crow + threshold, + ccol - threshold : ccol + threshold, + ] = scale + + if scales: + for scale_params in scales: + for scale_threshold, scale_value in scale_params: + scaled_scale_value = scale_value * strength + scale_mask = torch.ones(x_freq.shape, device=x.device) + scale_mask[ + ..., + crow - scale_threshold : crow + scale_threshold, + ccol - scale_threshold : ccol + scale_threshold, + ] = scaled_scale_value + mask = mask + (scale_mask - mask) * strength + + x_freq *= mask + + # IFFT + x_freq = fft.ifftshift(x_freq, dim=(-2, -1)) + return fft.ifftn(x_freq, dim=(-2, -1)).real.to(x.dtype) + + +class OpType(Enum): + # scale, strength, blend, blend mode, use hidden mean + SLICE = auto() + + # scale, filter, filter strength, threshold + FFILTER = auto() + + # type (bicubic, nearest, bilinear, area), scale width, scale height, antialias + SCALE_TORCH = auto() + + # type (bicubic, nearest, bilinear, area), antialias + UNSCALE_TORCH = auto() + + # type width (slerp, slerp_alt, hslerp, colorize), type height, scale width, scale height, antialias size + SCALE = auto() + + # type width (slerp, slerp_alt, hslerp, colorize), type height, antialias size + UNSCALE = auto() + + # direction (h, v) + FLIP = auto() + + # count + ROT90 = auto() + + # true/false - only makes sense with output block + TARGET_SKIP = auto() + + # factor + MULTIPLY = auto() + + +class CondType(Enum): + TYPE = auto() + BLOCK = auto() + STAGE = auto() + FROM_PERCENT = auto() + TO_PERCENT = auto() + PERCENT = auto() + COUNT = auto() + + +class BlockType(Enum): + INPUT = auto() + INPUT_AFTER_SKIP = auto() + MIDDLE = auto() + OUTPUT = auto() + + +class BlockCond: + def __init__(self, typ, value): + self.typ = getattr(CondType, typ.upper().strip()) + self.value = set(value if isinstance(value, (list, tuple)) else (value,)) + + def test(self, state, count=0): + match self.typ: + case CondType.FROM_PERCENT: + pct = state[CondType.PERCENT] + result = all(pct >= v for v in self.value) + case CondType.TO_PERCENT: + pct = state[CondType.PERCENT] + result = all(pct <= v for v in self.value) + case CondType.COUNT: + result = all(count < v for v in self.value) + case _: + result = state[self.typ] in self.value + return result + + def __repr__(self): + return f"" + + +class BlockConds: + def __init__(self, conds, count=0): + self.count = count + if not conds: + self.conds = () + return + self.conds = tuple(BlockCond(ct, cv) for ct, cv in conds.items()) + + def test(self, state): + count = self.count + result = all(c.test(state, count) for c in self.conds) + if result: + self.count += 1 + return result + + def __repr__(self): + return f"" + + +class BlockRule: + @classmethod + def from_dict(cls, val): + if not isinstance(val, (list, tuple)): + val = (val,) + + return tuple( + cls( + conds=d.get("if", ()), + ops=d.get("ops", ()), + matched=d.get("then", ()), + nomatched=d.get("else", ()), + ) + for d in val + ) + + def __init__(self, conds=(), ops=(), matched=(), nomatched=()): + self.conds = BlockConds(conds) + self.ops = tuple(BlockOp(o[0], *o[1:]) for o in ops) + self.matched = BlockRule.from_dict(matched) + self.nomatched = BlockRule.from_dict(nomatched) + + def get_all_types(self): + result = {c.value for c in self.conds if c.typ == CondType.TYPE} + for r in self.matched: + result |= r.get_all_types() + for r in self.nomatched: + result |= r.get_all_types() + return result + + def eval(self, state): + # print("EVAL", state | {"h": None, "hsp": None}) + if not self.conds.test(state): + for r in self.nomatched: + r.eval(state) + return + state["target"] = "h" + for o in self.ops: + o.eval(state) + for r in self.matched: + r.eval(state) + + def __repr__(self): + return f"" + + +class BlehBlockOps: + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + CATEGORY = "bleh/model_patches" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "rules": ("STRING", {"multiline": True, "forceInput": True}), + }, + } + + def patch( + self, + model, + rules, + ): + rules = rules.strip() + print("PARSING", repr(rules)) + if len(rules) == 0: + return (model.clone(),) + parsed_rules = yaml.safe_load(rules) + rules = tuple(BlockRule.from_dict(r)[0] for r in parsed_rules) + print("RULES:", rules) + + # Arbitrary number that should have good enough precision + pct_steps = 400 + pct_incr = 1.0 / pct_steps + sig2pct = tuple( + model.model.model_sampling.percent_to_sigma(x / pct_steps) + for x in range(pct_steps, -1, -1) + ) + + def get_pct(topts): + sigma = topts["sigmas"][0].item() + # This is obviously terrible but I couldn't find a better way to get the percentage from the current sigma. + idx = bisect.bisect_right(sig2pct, sigma) + if idx >= len(sig2pct): + # Sigma out of range somehow? + return None + return pct_incr * (pct_steps - idx) + + stages = (1280, 640, 320) + + def make_state(typ, topts, h, hsp=None): + pct = get_pct(topts) + if pct is None: + return None + + stage = stages.index(h.shape[1]) + 1 if h.shape[1] in stages else -1 + return { + CondType.TYPE: typ, + CondType.PERCENT: pct, + CondType.BLOCK: topts["block"][1], + CondType.STAGE: stage, + "h": h, + "hsp": hsp, + "target": "h", + } + + def block_patch(typ, h, topts): + state = make_state(typ, topts, h) + if state is None: + return h + for rule in rules: + rule.eval(state) + return state["h"] + + def output_block_patch(h, hsp, transformer_options): + state = make_state("output", transformer_options, h, hsp) + if state is None: + return h + for rule in rules: + rule.eval(state) + return state["h"], state["hsp"] + + m = model.clone() + m.set_model_input_block_patch_after_skip( + lambda *args: block_patch("input_after_skip", *args), + ) + m.set_model_input_block_patch(lambda *args: block_patch("input", *args)) + m.set_model_patch( + lambda *args: block_patch("middle", *args), + "middle_block_patch", + ) + m.set_model_output_block_patch(output_block_patch) + return (m,) + + +class BlehLatentUpscaleBy: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "samples": ("LATENT",), + "upscale_method": (UPSCALE_METHODS,), + "scale_width": ( + "FLOAT", + {"default": 1.5, "min": 0.01, "max": 8.0, "step": 0.01}, + ), + "scale_height": ( + "FLOAT", + {"default": 1.5, "min": 0.01, "max": 8.0, "step": 0.01}, + ), + "antialias": ("BOOLEAN", {"default": False}), + }, + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "upscale" + + CATEGORY = "latent" + + def upscale(self, samples, upscale_method, scale_width, scale_height): + s = samples.copy() + width = round(samples["samples"].shape[3] * scale_width) + height = round(samples["samples"].shape[2] * scale_height) + s["samples"] = scale_samples( + samples["samples"], + width, + height, + upscale_method, + ) + return (s,) diff --git a/py/nodes/deepShrink.py b/py/nodes/deepShrink.py new file mode 100644 index 0000000..b8b0b95 --- /dev/null +++ b/py/nodes/deepShrink.py @@ -0,0 +1,171 @@ +# Adapted from the ComfyUI built-in node + +import bisect + +import torch +from comfy.utils import bislerp + + +class DeepShrinkBleh: + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + CATEGORY = "bleh/model_patches" + + upscale_methods = ( + "bicubic", + "nearest-exact", + "bilinear", + "area", + "bislerp", + ) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "commasep_block_numbers": ( + "STRING", + { + "default": "3", + }, + ), + "downscale_factor": ( + "FLOAT", + {"default": 2.0, "min": 1.0, "max": 32.0, "step": 0.1}, + ), + "start_percent": ( + "FLOAT", + {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}, + ), + "start_fadeout_percent": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}, + ), + "end_percent": ( + "FLOAT", + {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.001}, + ), + "downscale_after_skip": ("BOOLEAN", {"default": True}), + "downscale_method": (cls.upscale_methods,), + "upscale_method": (cls.upscale_methods,), + "antialias_downscale": ("BOOLEAN", {"default": False}), + "antialias_upscale": ("BOOLEAN", {"default": False}), + }, + } + + def patch( + self, + model, + commasep_block_numbers, + downscale_factor, + start_percent, + start_fadeout_percent, + end_percent, + downscale_after_skip, + downscale_method, + upscale_method, + antialias_downscale, + antialias_upscale, + ): + block_numbers = tuple( + int(x) for x in commasep_block_numbers.split(",") if x.strip() + ) + downscale_factor = 1.0 / downscale_factor + if not (block_numbers and all(val > 0 and val <= 32 for val in block_numbers)): + raise ValueError( + "BlehDeepShrink: Bad value for block numbers: must be comma-separated list of numbers between 1-32", + ) + antialias_downscale = antialias_downscale and downscale_method in ( + "bicubic", + "bilinear", + ) + antialias_upscale = antialias_upscale and upscale_method in ( + "bicubic", + "bilinear", + ) + if start_fadeout_percent < start_percent: + start_fadeout_percent = start_percent + elif start_fadeout_percent > end_percent: + # No fadeout. + start_fadeout_percent = 1000.0 + + sigma_start = model.model.model_sampling.percent_to_sigma(start_percent) + sigma_end = model.model.model_sampling.percent_to_sigma(end_percent) + # Arbitrary number that should have good enough precision + pct_steps = 400 + pct_incr = 1.0 / pct_steps + sig2pct = tuple( + model.model.model_sampling.percent_to_sigma(x / pct_steps) + for x in range(pct_steps, -1, -1) + ) + + def input_block_patch(h, transformer_options): + sigma = transformer_options["sigmas"][0].item() + if ( + sigma > sigma_start + or sigma < sigma_end + or transformer_options["block"][1] not in block_numbers + ): + return h + # This is obviously terrible but I couldn't find a better way to get the percentage from the current sigma. + idx = bisect.bisect_right(sig2pct, sigma) + if idx >= len(sig2pct): + # Sigma out of range somehow? + return h + pct = pct_incr * (pct_steps - idx) + if ( + pct < start_fadeout_percent + or start_fadeout_percent > end_percent + or pct > end_percent + ): + scaled_scale = downscale_factor + else: + # May or not be accurate but the idea is to scale the downscale factor by the percentage + # of the start fade to end deep shrink we have currently traversed. It at least sort of works. + downscale_pct = 1.0 - ( + (pct - start_fadeout_percent) + / (end_percent - start_fadeout_percent) + ) + scaled_scale = 1.0 - ((1.0 - downscale_factor) * downscale_pct) + orig_width, orig_height = h.shape[-1], h.shape[-2] + width, height = ( + round(orig_width * scaled_scale), + round(orig_height * scaled_scale), + ) + if scaled_scale >= 0.98 or width >= orig_width or height >= orig_height: + return h + if downscale_method == "bislerp": + return bislerp(h, width, height) + return torch.nn.functional.interpolate( + h, + size=(height, width), + mode=downscale_method, + antialias=antialias_downscale, + ) + + def output_block_patch(h, hsp, _transformer_options): + if h.shape[2] == hsp.shape[2]: + return h, hsp + if upscale_method == "bislerp": + return bislerp( + h, + hsp.shape[-1], + hsp.shape[-2], + ), hsp + return torch.nn.functional.interpolate( + h, + size=(hsp.shape[-2], hsp.shape[-1]), + mode=upscale_method, + antialias=antialias_upscale, + ), hsp + + m = model.clone() + if downscale_factor == 0.0 or start_percent >= 1.0: + return (m,) + if downscale_after_skip: + m.set_model_input_block_patch_after_skip(input_block_patch) + else: + m.set_model_input_block_patch(input_block_patch) + m.set_model_output_block_patch(output_block_patch) + return (m,) diff --git a/py/hypertile.py b/py/nodes/hyperTile.py similarity index 100% rename from py/hypertile.py rename to py/nodes/hyperTile.py diff --git a/py/modelPatchConditional.py b/py/nodes/modelPatchConditional.py similarity index 100% rename from py/modelPatchConditional.py rename to py/nodes/modelPatchConditional.py diff --git a/py/refinerAfter.py b/py/nodes/refinerAfter.py similarity index 100% rename from py/refinerAfter.py rename to py/nodes/refinerAfter.py diff --git a/py/samplers.py b/py/nodes/samplers.py similarity index 100% rename from py/samplers.py rename to py/nodes/samplers.py diff --git a/py/sigmas.py b/py/nodes/sigmas.py similarity index 100% rename from py/sigmas.py rename to py/nodes/sigmas.py diff --git a/ruff.toml b/ruff.toml index 9a26a38..4668c97 100644 --- a/ruff.toml +++ b/ruff.toml @@ -13,6 +13,7 @@ ignore = [ "D102", "D103", "D104", + "D105", "D107", "D211", "D213",