Add many new scaling types + enhancement effects (#4)

* Add many new scaling types * Add enhancement effects
2024-05-06 16:48:47 -06:00
parent ddfad286d6
commit ad0ac75602
6 changed files with 237 additions and 37 deletions
@@ -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!
@@ -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"]
@@ -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).
@@ -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(
@@ -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()
@@ -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,