From ad0ac75602f5082733a11cdc611c3ed4810fbdc6 Mon Sep 17 00:00:00 2001 From: blepping <157360029+blepping@users.noreply.github.com> Date: Mon, 6 May 2024 16:48:47 -0600 Subject: [PATCH] Add many new scaling types + enhancement effects (#4) * Add many new scaling types * Add enhancement effects --- README.md | 74 ++++++++++++++---- __init__.py | 4 +- changelog.md | 5 ++ py/latent_utils.py | 171 ++++++++++++++++++++++++++++++++++++----- py/nodes/deepShrink.py | 9 ++- py/nodes/ops.py | 11 +++ 6 files changed, 237 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index a11fe6b..2e7e45d 100644 --- a/README.md +++ b/README.md @@ -367,6 +367,12 @@ These two shortcuts can be combined. A mask of `[["rep", 2, 1, [3, 0], 2]]` expa ] ``` +**`apply_enhancement`**: Applies an [enhancement](#enhancement-types) to the target. + +1. `scale`: 1.0 +2. `type`: korniabilateralblur + + #### Blend Modes 1. bislerp: Interpolates between tensors a and b using normalized linear interpolation. @@ -400,18 +406,7 @@ Custom filters may also be defined. For example, `gaussianblur` in the YAML filt #### Scaling Functions -1. bicubic: Generally the best option. -2. bilinear -3. nearest-exact -4. area -5. bislerp: Interpolates between tensors a and b using normalized linear interpolation. -6. colorize: Supposedly transfers color. May or may not work that way. -7. hslerp: Hybrid Spherical Linear Interporation, supposedly smooths transitions between orientations and colors. -8. bibislerp: Uses bislerp as the slerp function in bislerp. When slerping once just isn't enough. -9. cosinterp: Cosine interpolation. -10. cuberp: Cubic interpolation. -11. inject: Adds the value scaled by the ratio. Probably not the best for scaling. -12. lineardodge: Supposedly simulates a brightning effect. +See [Scaling Types](#scaling-types) below. #### Examples @@ -453,9 +448,62 @@ Custom filters may also be defined. For example, `gaussianblur` in the YAML filt ### BlehLatentOps -Basically the same as BlehBlockOps, except the condition `type` will be `latent`. Obviously stuff involving steps, percentages, etc does not apply. +Basically the same as BlehBlockOps, except the condition `type` will be `latent`. Obviously stuff involving steps, percentages, etc do not apply. This node allows you to apply the blending/filtering/scaling operations to a latent. +### BlehLatentScaleBy + +Like the builtin `LatentScaleBy` node, however it allows setting the horizontal and vertical scaling types and scales independently +as well as allowing providing an extended list of scaling options. Can also be useful for testing what different types of scaling or +enhancement effects look like. + +## Scaling Types + +* bicubic: Generally the safe option. +* bilinear: Like bicubic but slightly not as good? +* nearest-exact +* area +* bislerp: Interpolates between tensors a and b using normalized linear interpolation. +* colorize: Supposedly transfers color. May or may not work that way. +* hslerp: Hybrid Spherical Linear Interporation, supposedly smooths transitions between orientations and colors. +* bibislerp: Uses bislerp as the slerp function in bislerp. When slerping once just isn't enough. +* cosinterp: Cosine interpolation. +* cuberp: Cubic interpolation. +* inject: Adds the value scaled by the ratio. Probably not the best for scaling. +* lineardodge: Supposedly simulates a brightning effect. +* random: Chooses a random relatively normal scaling function each time. My thought is this will avoid artifacts from + a specific scaling type from getting reinforced each step. Generally only useful for Deep Shrink or + [jankhdiffusion](https://github.com/blepping/comfyui_jankhidiffusion). +* randomaa: Like `random`, however it will also choose a random antialias size. + +Scaling types like `bicubic+something` will apply the `something` enhancement after scaling. See below. + +Scaling types that start with `rev` like `revinject` reverse the arguments to the scaling function. +For example, `inject` does `a + b * scale`, `revinject` does `b + a * scale`. When is this desirable? +I really don't know! Just stuff to experiment with. It may or may not be useful. (`revcosinterp` looks better than `cosinterp` though.) + +**Note**: Scaling types like `random` are very experimental and may be modified or removed. + +## Enhancement Types + +* randmultihighlowpass: Randomly uses multihighpass or multilowpass filter. Effect is generally quite strong. +* randhilowpass: Randomly uses a highpass or lowpass filter. When you filter both high and low frequencies you are left with... + nothing! The effect is very strong. May not be useful. +* randlowbandpass: Randomly uses a bandpass or lowpass filter. +* randhibandpass: Randomly uses a bandpass or highpass filter. +* renoise1: Adds some gaussian noise. Starts off relatively weak and increases based on sigma. +* renoise2: Adds some guassian noise. Starts relatively strong and decreases based on sigma. +* korniabilateralblur: Applies a bilateral (edge preserving) blur effect. +* korniagaussianblur: Applies a guassian blur effect. +* korniasharpen: Applies a sharpen effect. +* korniaedge: Applies an edge enhancement effect. +* korniarevedge: Applies an edge softening effect - may not work correctly. +* korniarandblursharp: Randomly chooses between blurring and sharpening. + +Also may be an item from [Filters](#filters). + +**Note**: These enhancements are very experimental and may be modified or removed. + ## Credits Latent blending and scaling and filter functions based on implementation from https://github.com/WASasquatch/FreeU_Advanced - thanks! diff --git a/__init__.py b/__init__.py index 5eae490..4b7a47e 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,7 @@ from .py import settings +BLEH_VERSION = 0 + settings.load_settings() if settings.SETTINGS.btp_enabled: @@ -35,4 +37,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "BlehDeepShrink": "Kohya Deep Shrink (bleh)", } -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "BLEH_VERSION"] diff --git a/changelog.md b/changelog.md index c45f80b..8231eb2 100644 --- a/changelog.md +++ b/changelog.md @@ -2,6 +2,11 @@ Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top. +## 20240506 + +* Add many new scaling types. +* Add enhancements that can be combined with scaling, also `apply_enhancement` blockops function. + ## 20240423 * Added `BlehPlug` and `BlehDisableNoise` (see README for usage and description). diff --git a/py/latent_utils.py b/py/latent_utils.py index e5a57b7..d22d3d0 100644 --- a/py/latent_utils.py +++ b/py/latent_utils.py @@ -3,6 +3,7 @@ import math +import kornia.filters as kf import numpy as np import torch from torch import fft @@ -126,6 +127,10 @@ BLENDING_MODES = { # Simulates a brightening effect by adding tensor b to tensor a, scaled by t. "lineardodge": lambda a, b, t: normalize(a + b * t), } +for k in tuple(BLENDING_MODES.keys()): + if k == "hslerp": + continue + BLENDING_MODES[f"rev{k}"] = lambda a, b, t, f=BLENDING_MODES[k]: f(b, a, t) FILTER_PRESETS = { "none": (), @@ -177,20 +182,59 @@ FILTER_PRESETS = { "multisharpen": ((5, 1.5), (10, 2.0), (15, 2.5)), # Multi-scale sharpening } + +BIDERP_MODES = {k: v for k, v in BLENDING_MODES.items() if not k.endswith("slerp")} +BIDERP_MODES |= { + "hslerp": hslerp_alt, + "bislerp": slerp_orig, + "bibislerp": BLENDING_MODES["bislerp"], + "revhslerp": lambda a, b, t, f=hslerp_alt: f(b, a, t), + "revbislerp": lambda a, b, t, f=slerp_orig: f(b, a, t), + "revbibislerp": BLENDING_MODES["revbislerp"], +} + + UPSCALE_METHODS = ( "bicubic", "nearest-exact", "bilinear", "area", - "bislerp", + *BIDERP_MODES.keys(), + *( + f"{meth}+{enh}" + for meth in ("bicubic", "bislerp", "hslerp", "random") + for enh in ( + "lowpass", + "highpass", + "bandpass", + "randhilowpass", + "randmultihilowpass", + "randhibandpass", + "randlowbandpass", + "gaussianblur", + "edge", + "sharpen", + "korniabilateralblur", + "korniagaussianblur", + "korniasharpen", + "korniaedge", + "korniarevedge", + "korniarandblursharp", + "renoise1", + "renoise2", + ) + ), + "random", + "randomaa", +) + + +RAND_UPSCALE_METHODS = ( + "bicubic", "colorize", - "hslerp", - "bibislerp", - "cosinterp", - "cuberp", - "inject", - "lerp", - "lineardodge", + "bislerp", + "revcosinterp", + "bilinear", ) FILTER_SIZES = ( @@ -217,6 +261,75 @@ def antialias_tensor(x, antialias_size): return torch.nn.functional.conv2d(x, filt, groups=channels, padding="same") +def enhance_tensor(x, name, scale=1.0, sigma=None): # noqa: PLR0911 + randitems = None + match name: + case "randmultihilowpass": + scale *= 0.1 + randskip = 4 + randitems = ("multilowpass", "multihighpass") + case "randhilowpass": + scale *= 0.1 + randskip = 6 + randitems = ("lowpass", "highpass") + case "randlowbandpass": + scale *= 0.25 + randskip = 1 + randitems = ("lowpass", "multilowpass", "bandpass") + case "randhibandpass": + scale *= 0.25 + randskip = 1 + randitems = ("highpass", "multihighpass", "bandpass") + case "bandpass": + scale *= 0.2 + case "renoise1" | "renoise2": + if sigma is None: + return x + noise_scale = ( + min(sigma / 6.0, 2.0 / max(sigma, 1e-05)) + if name == "renoise1" + else sigma / 8.0 + ) + if noise_scale < 1e-04: + return x + noise = torch.randn_like(x) + return noise.mul_(noise_scale).add_(x) + if randitems: + ridx = torch.randint(len(randitems) + randskip, (1,), device="cpu").item() + if ridx >= len(randitems): + return x + return enhance_tensor(x, randitems[ridx], scale=scale) + fpreset = FILTER_PRESETS.get(name) + if fpreset is not None: + return ffilter(x, 1, 1.0, fpreset, 0.5 * scale) + match name: + case "korniabilateralblur": + return x + (kf.bilateral_blur(x, (3, 3), 0.1, (1.5, 1.5)) - x) * ( + scale * 2.0 + ) + case "korniagaussianblur": + return kf.gaussian_blur2d(x, (3, 3), (1.5, 1.5)) * scale + case "korniasharpen": + return x + (kf.unsharp_mask(x, (3, 3), (1.5, 1.5)) - x) * (scale / 2.0) + case "korniaedge" | "korniarevedge": + blur = kf.bilateral_blur(x, (3, 3), 0.1, (1.5, 1.5)) - x + sharpened = kf.unsharp_mask(x, (3, 3), (1.5, 1.5)) - x + if name == "korniarevedge": + scale *= -1.0 + return x + (sharpened + blur) * (scale / 2.0) + case "korniarandblursharp": + return enhance_tensor( + x, + "korniagaussianblur" + if torch.rand(1, device="cpu").item() < 0.5 + else "korniasharpen", + scale=scale, + ) + case _: + raise ValueError("Unknown enhancement") + + +@torch.no_grad() def scale_samples( samples, width, @@ -224,9 +337,28 @@ def scale_samples( mode="bicubic", mode_h=None, antialias_size=0, + post_effect_strength=1.0, + sigma=None, ): if mode_h is None: mode_h = mode + mode, *enhancement = mode.split("+", 1) + mode_h = mode_h.split("+", 1)[0] + modes = (mode, mode_h) + if "randomaa" in modes: + raasize, useraa = torch.rand(2, device="cpu").detach() + antialias_size = (int(raasize * 7) + 1) * int(useraa * 2) + if "random" in modes or "randomaa" in modes: + ridxs = torch.randint( + len(RAND_UPSCALE_METHODS), + (2,), + dtype=torch.uint8, + ).tolist() + mode, mode_h = ( + m if mode not in ("random", "randomaa") else RAND_UPSCALE_METHODS[ridx] + for ridx, m in zip(ridxs, (mode, mode_h)) + ) + mode_h = mode if mode in ("bicubic", "nearest-exact", "bilinear", "area"): result = torch.nn.functional.interpolate( samples, @@ -236,6 +368,13 @@ def scale_samples( ) else: result = biderp(samples, width, height, mode, mode_h) + if enhancement: + result = enhance_tensor( + result, + enhancement[-1], + scale=post_effect_strength, + sigma=sigma, + ) if antialias_size < 1 or antialias_size > 7: return result return antialias_tensor(result, antialias_size) @@ -246,18 +385,10 @@ def biderp(samples, width, height, mode="bislerp", mode_h=None): if mode_h is None: mode_h = mode - modes = { - "colorize": BLENDING_MODES["colorize"], - "hslerp": hslerp_alt, - "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) + derp_w, derp_h = ( + BIDERP_MODES.get(mode, slerp_orig), + BIDERP_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( diff --git a/py/nodes/deepShrink.py b/py/nodes/deepShrink.py index a611e10..be12694 100644 --- a/py/nodes/deepShrink.py +++ b/py/nodes/deepShrink.py @@ -100,7 +100,7 @@ class DeepShrinkBleh: ) def input_block_patch(h, transformer_options): - sigma = transformer_options["sigmas"][0].item() + sigma = transformer_options["sigmas"][0].cpu().item() if ( sigma > sigma_start or sigma < sigma_end @@ -140,17 +140,20 @@ class DeepShrinkBleh: height, mode=downscale_method, antialias_size=3 if antialias_downscale else 0, + sigma=sigma, ) - def output_block_patch(h, hsp, _transformer_options): - if h.shape[2] == hsp.shape[2]: + def output_block_patch(h, hsp, transformer_options): + if h.shape[-2:] == hsp.shape[-2:]: return h, hsp + sigma = transformer_options["sigmas"][0].cpu().item() return latent_utils.scale_samples( h, hsp.shape[-1], hsp.shape[-2], mode=upscale_method, antialias_size=3 if antialias_upscale else 0, + sigma=sigma, ), hsp m = model.clone() diff --git a/py/nodes/ops.py b/py/nodes/ops.py index 213b37b..d1c04bc 100644 --- a/py/nodes/ops.py +++ b/py/nodes/ops.py @@ -122,6 +122,9 @@ class OpType(Enum): # count, [ops] REPEAT = auto() + # + APPLY_ENHANCEMENT = auto() + OP_DEFAULTS = { OpType.SLICE: OrderedDict( @@ -189,6 +192,7 @@ OP_DEFAULTS = { ), OpType.CROP: OrderedDict(top=0, bottom=0, left=0, right=0), OpType.REPEAT: OrderedDict(count=2, ops=()), + OpType.APPLY_ENHANCEMENT: OrderedDict(scale=1.0, type="korniabilateralblur"), } @@ -410,6 +414,7 @@ class OpScaleTorch(Operation): height, mode, antialias_size=8 if antialias else 0, + sigma=state.get("sigma"), ) @@ -657,6 +662,12 @@ class OpRepeat(SubOpsOperation): return state[state["target"]] +class OpApplyEnhancement(Operation): + def op(self, t, _state): + scale, typ = self.args + return enhance_tensor(t, typ, scale=scale, sigma=state.get("sigma")) + + OP_TO_OPCLASS = { OpType.SLICE: OpSlice, OpType.FFILTER: OpFFilter,