SD.Next extension for additional upscalers based on chaiNNer
#!/usr/bin/env python

import os
import sys
import time
import warnings
import torch
import numpy as np
from PIL import Image
from nodes.impl.upscale.tiler import MaxTileSize, NoTiling, Tiler
from nodes.impl.pytorch.auto_split import pytorch_auto_split
from nodes.impl.pytorch.types import PyTorchSRModel
from nodes.load_model import load_model
from nodes.log import logger


warnings.filterwarnings('ignore', category=UserWarning) # disable those for now as many backends reports tons
fp16 = False # HAT does not support fp16
device = torch.device('cuda')


def parse_tile_size_input(tile_size: int) -> Tiler:
    if tile_size == 0:
        return MaxTileSize(tile_size)
    elif tile_size == -1:
        return NoTiling()
    elif tile_size == -2:
        return MaxTileSize()
    elif tile_size < 0:
        raise ValueError(f"ChaiNNer invalid tile size: {tile_size}")
    return MaxTileSize(tile_size)


def upscale(image: Image, model: PyTorchSRModel, tile: int = 256):
    img = np.array(image)
    with torch.no_grad():
        upscaled = pytorch_auto_split(img, model=model, device=device, use_fp16=fp16, tiler=parse_tile_size_input(tile))
        return Image.fromarray(np.uint8(256 * upscaled))


if __name__ == "__main__":
    sys.argv.pop(0)
    if len(sys.argv) == 0:
        logger.error('chainner: no files specified')
        sys.exit(1)
    for modelfile in os.listdir('models'):
        try:
            modelname = os.path.splitext(modelfile)[0]
            srmodel: PyTorchSRModel = load_model(os.path.join("models", modelfile), device=device, fp16=fp16)
            logger.info(f'model="{modelname}" arch="{srmodel.__class__.__name__}" scale={srmodel.scale}')
            for imagename in sys.argv:
                if not os.path.isfile(imagename):
                    logger.error(f'image={imagename} not found')
                    continue
                inputimage = Image.open(imagename).convert('RGB')
                t0 = time.time()
                outputimage = upscale(image=inputimage, model=srmodel, tile=256)
                t1 = time.time()
                base, ext = os.path.splitext(imagename)
                outputname = f'{base}-{modelname}{ext}'
                outputimage.save(outputname)
                logger.info(f'input="{imagename}" {inputimage.size} output="{outputname}" {outputimage.size} time={t1-t0:.2f}s')
            srmodel = None
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
            # sys.exit(1)
        except Exception as e:
            logger.error(f'Error: fn={modelfile} {e}')

chainner-test.py
vladmandic/sd-extension-chainner
from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import List, Literal, Optional, Type, TypedDict, Union, TYPE_CHECKING

from base_types import InputId

if TYPE_CHECKING:
    import navi

InputKind = Literal[
    "number",
    "slider",
    "dropdown",
    "text",
    "directory",
    "file",
    "color",
    "generic",
]


@dataclass
class InputConversion:
    """
    An input conversion can be used to convert the assigned type of an input.
    This is useful to model the changes `enforce` makes to values.

    `type` is used to declare which type is intended to be converted by this
    conversion. `convert` is the expression that does the actual conversion. It
    will be given a special parameter called `Input` that will be the value to
    convert. The `Input` parameter is guaranteed to be a non-empty sub type of
    `type`.

    Example:
    To convert all numbers to string, use this conversions:
    ```
    InputConversion("number", "toString(Input)")
    ```
    """

    type: navi.ExpressionJson
    convert: navi.ExpressionJson

    def toDict(self):
        return {
            "type": self.type,
            "convert": self.convert,
        }


class LiteralErrorValue(TypedDict):
    type: Literal["literal"]
    value: str | int | float | None


class FormattedErrorValue(TypedDict):
    type: Literal["formatted"]
    formatString: str


class UnknownErrorValue(TypedDict):
    type: Literal["unknown"]
    typeName: str
    typeModule: str


ErrorValue = Union[LiteralErrorValue, FormattedErrorValue, UnknownErrorValue]


class BaseInput:
    def __init__(
        self,
        input_type: navi.ExpressionJson,
        label: str,
        kind: InputKind = "generic",
        has_handle=True,
        associated_type: Union[Type, None] = None,
    ):
        self.input_type: navi.ExpressionJson = input_type
        self.input_conversions: List[InputConversion] = []
        self.input_adapt: navi.ExpressionJson | None = None
        self.type_definitions: str | None = None
        self.kind: InputKind = kind
        self.label: str = label
        self.optional: bool = False
        self.has_handle: bool = has_handle
        self.id: InputId = InputId(-1)
        self.associated_type: Type = associated_type

        # Optional documentation
        self.description: str | None = None
        self.hint: bool = False

    # This is the method that should be created by each input
    def enforce(self, value: object):
        """Enforce the input type"""
        return value

    # This is the method that should be called by the processing code
    def enforce_(self, value: object | None):
        if self.optional and value is None:
            return None
        assert value is not None, (
            f"Expected value to exist, "
            f"but does not exist for {self.kind} input with type {self.input_type} and label {self.label}"
        )
        return self.enforce(value)

    def get_error_value(self, value: object) -> ErrorValue:
        if isinstance(value, Enum):
            # unwrap enum
            value = value.value

        if isinstance(value, bool):
            # bools need to be 0 or 1
            return {"type": "literal", "value": int(value)}

        if isinstance(value, (int, float, str)) or value is None:
            return {"type": "literal", "value": value}

        return {
            "type": "unknown",
            "typeName": type(value).__qualname__,
            "typeModule": type(value).__module__,
        }

    def toDict(self):
        actual_type = [self.input_type, "null"] if self.optional else self.input_type
        return {
            "id": self.id,
            "type": actual_type,
            "conversions": [c.toDict() for c in self.input_conversions],
            "adapt": self.input_adapt,
            "typeDefinitions": self.type_definitions,
            "kind": self.kind,
            "label": self.label,
            "optional": self.optional,
            "hasHandle": self.has_handle,
            "description": self.description,
            "hint": self.hint,
        }

    def with_id(self, input_id: InputId | int):
        self.id = InputId(input_id)
        return self

    def with_docs(self, *description: str, hint=False):
        self.description = "\n\n".join(description)
        self.hint = hint
        return self

    def make_optional(self):
        self.optional = True
        if self.associated_type is not None:
            associated_type = self.associated_type
            self.associated_type = Optional[associated_type]
        return self

    def __repr__(self):
        return str(self.toDict())

    def __iter__(self):
        yield from self.toDict().items()

nodes/base_input.py
vladmandic/sd-extension-chainner
from __future__ import annotations

from typing import Literal, Type, Union, TYPE_CHECKING

from base_types import OutputId

if TYPE_CHECKING:
    import navi

OutputKind = Literal["image", "large-image", "tagged", "generic"]


class BaseOutput:
    def __init__(
        self,
        output_type: navi.ExpressionJson,
        label: str,
        kind: OutputKind = "generic",
        has_handle: bool = True,
        associated_type: Union[Type, None] = None,
    ):
        self.output_type: navi.ExpressionJson = output_type
        self.label: str = label
        self.id: OutputId = OutputId(-1)
        self.never_reason: str | None = None
        self.kind: OutputKind = kind
        self.has_handle: bool = has_handle

        self.associated_type: Union[Type, None] = associated_type

        # Optional documentation
        self.description: str | None = None

    def toDict(self):
        return {
            "id": self.id,
            "type": self.output_type,
            "label": self.label,
            "neverReason": self.never_reason,
            "kind": self.kind,
            "hasHandle": self.has_handle,
            "description": self.description,
        }

    def with_id(self, output_id: OutputId | int):
        self.id = OutputId(output_id)
        return self

    def with_never_reason(self, reason: str):
        self.never_reason = reason
        return self

    def with_docs(self, *description: str):
        self.description = "\n\n".join(description)
        return self

    def __repr__(self):
        return str(self.toDict())

    def __iter__(self):
        yield from self.toDict().items()

    def get_broadcast_data(self, _value):
        return None

    def get_broadcast_type(self, _value) -> navi.ExpressionJson | None:
        return None

    def enforce(self, value: object) -> object:
        assert value is not None
        return value

nodes/base_output.py
vladmandic/sd-extension-chainner
from typing import Any, Dict, Generic, List, NewType, Optional, TypeVar, Union

from base_types import InputId
from nodes.base_input import BaseInput

T = TypeVar("T")


GroupId = NewType("GroupId", int)


class GroupInfo:
    def __init__(
        self,
        group_id: GroupId,
        kind: str,
        options: Optional[Dict[str, Any]] = None,
    ) -> None:
        self.id: GroupId = group_id
        self.kind: str = kind
        self.options: Dict[str, Any] = {} if options is None else options


class Group(Generic[T]):
    def __init__(self, info: GroupInfo, items: List[T]) -> None:
        self.info: GroupInfo = info
        self.items: List[T] = items

    def toDict(self):
        return {
            "id": self.info.id,
            "kind": self.info.kind,
            "options": self.info.options,
            "items": [i.toDict() if isinstance(i, Group) else i for i in self.items],
        }


NestedGroup = Group[Union[BaseInput, "NestedGroup"]]
NestedIdGroup = Group[Union[InputId, "NestedIdGroup"]]


# pylint: disable-next=redefined-builtin
def group(kind: str, options: Optional[Dict[str, Any]] = None, id: int = -1):
    info = GroupInfo(GroupId(id), kind, options)

    def ret(*items: Union[BaseInput, NestedGroup]) -> NestedGroup:
        return Group(info, list(items))

    return ret

nodes/group.py
vladmandic/sd-extension-chainner
from __future__ import annotations

from enum import Enum
from typing import Iterable, List, Literal, Tuple, TypedDict, Union, TYPE_CHECKING

from base_types import InputId
from nodes.group import NestedGroup, group

if TYPE_CHECKING:
    import navi
    from nodes.base_input import BaseInput

InputValue = Union[int, str]
EnumValues = Union[
    InputValue,
    Enum,
    Iterable[str],
    Iterable[int],
    Iterable[Enum],
]
RawEnumValues = Union[
    InputValue, List[str], List[int], Tuple[str, ...], Tuple[int, ...]
]

_ConditionJson = Union[
    "_AndConditionJson",
    "_OrConditionJson",
    "_NotConditionJson",
    "_EnumConditionJson",
    "_TypeConditionJson",
]


class _AndConditionJson(TypedDict):
    kind: Literal["and"]
    items: List[_ConditionJson]


class _OrConditionJson(TypedDict):
    kind: Literal["or"]
    items: List[_ConditionJson]


class _NotConditionJson(TypedDict):
    kind: Literal["not"]
    condition: _ConditionJson


class _EnumConditionJson(TypedDict):
    kind: Literal["enum"]
    enum: InputId
    values: List[str | int]


class _TypeConditionJson(TypedDict):
    kind: Literal["type"]
    input: InputId
    condition: navi.ExpressionJson
    ifNotConnected: bool


class Condition:
    def __init__(self, value: _ConditionJson) -> None:
        self._value: _ConditionJson = value

    def to_json(self):
        return self._value

    def __and__(self, other: Condition) -> Condition:
        return Condition({"kind": "and", "items": [self._value, other._value]})

    def __or__(self, other: Condition) -> Condition:
        return Condition({"kind": "or", "items": [self._value, other._value]})

    def __invert__(self) -> Condition:
        return Condition({"kind": "not", "condition": self._value})

    @staticmethod
    def enum(enum: int, values: EnumValues) -> Condition:
        """
        A condition to check whether a certain dropdown/enum input has a certain value.
        """

        v: List[str | int] = []

        def convert(value: int | str | Enum):
            if isinstance(value, (int, str)):
                v.append(value)
            else:
                enum_value = value.value
                assert isinstance(enum_value, (int, str))
                v.append(enum_value)

        if isinstance(values, (int, str, Enum)):
            convert(values)
        else:
            for value in values:
                convert(value)

        return Condition(
            {
                "kind": "enum",
                "enum": InputId(enum),
                "values": v,
            }
        )

    @staticmethod
    def bool(input_id: int, value: bool) -> Condition:
        """
        A condition to check whether a certain bool input has a certain value.
        """
        return Condition(
            {
                "kind": "enum",
                "enum": InputId(input_id),
                "values": [int(value)],
            }
        )

    @staticmethod
    def type(
        input_id: int,
        condition: navi.ExpressionJson,
        if_not_connected: bool = False,
    ) -> Condition:
        """
        A condition to check whether a certain input is compatible a certain type.
        Here "compatible" is defined as overlapping.
        """
        return Condition(
            {
                "kind": "type",
                "input": InputId(input_id),
                "condition": condition,
                "ifNotConnected": if_not_connected,
            }
        )

    @staticmethod
    def const(value: bool) -> Condition:
        if value:
            return Condition({"kind": "and", "items": []})
        return Condition({"kind": "or", "items": []})


def if_group(condition: Condition):
    return group("conditional", {"condition": condition.to_json()})


def if_enum_group(enum: int, condition: EnumValues):
    return if_group(Condition.enum(enum, condition))


def required(condition: Condition | None = None):
    """
    Given generic inputs (meaning of kind "generic") that are optional, this group marks them as
    being required under the given condition. If no condition is given, `True` will be used.

    In addition to the given condition, if the require group is nested within conditional group
    (`if_group` and derivatives), then the conditions of all ancestor conditional groups must also
    be met.

    Note that this group only guarantees **best effort**. It cannot guarantee that the optional
    input is going to have a value if the condition is met. You must always check `None`.

    Example:
    ```py
    if_group(someCondition)(
        required()(
            GenericInput("Foo").make_optional(),
        )
    )
    ```

    In this example, the input "Foo" is required if and only if the input is visible (by virtue of
    the parent conditional group).
    """

    if condition is None:
        condition = Condition.const(True)
    return group("required", {"condition": condition.to_json()})


def seed_group(seed_input: BaseInput):
    """
    This groups is a wrapper around the `SeedInput`. It changes its visual appearance and adds a
    little button for users to click on to generate a new seed.

    All `SeedInput`s must be wrapped in this group.

    Example:
    ```py
    seed_group(SeedInput())
    ```
    """
    return group("seed")(seed_input)


def optional_list_group(*inputs: BaseInput | NestedGroup):
    """
    This groups wraps around optional inputs and displays them as a list.

    This can be used to create nodes that have a variable number of inputs. The user will initially
    see no inputs, but can add as many inputs as the group contains. While not true varargs, this
    can be used to create a similar effect.

    See the Text Append node for an example.
    """
    return group("optional-list")(*inputs)


def linked_inputs_group(*inputs: BaseInput):
    """
    This group wraps around inputs of the same type. It ensures that all inputs have the same
    value.

    "The same type" here not only refers to the Navi type of those inputs. All possible values
    from all inputs must also be valid values for all other inputs. This typically necessitates
    that the inputs are of the same class and use the same parameters.
    """
    return group("linked-inputs")(*inputs)

nodes/groups.py
vladmandic/sd-extension-chainner
from enum import Enum

import cv2
import numpy as np

from ..utils.utils import get_h_w_c
from .image_utils import as_target_channels, normalize, to_uint8


class BlendMode(Enum):
    NORMAL = 0
    DARKEN = 2
    MULTIPLY = 1
    COLOR_BURN = 5
    LINEAR_BURN = 22
    LIGHTEN = 3
    SCREEN = 12
    COLOR_DODGE = 6
    ADD = 4
    OVERLAY = 9
    SOFT_LIGHT = 17
    HARD_LIGHT = 18
    VIVID_LIGHT = 19
    LINEAR_LIGHT = 20
    PIN_LIGHT = 21
    REFLECT = 7
    GLOW = 8
    DIFFERENCE = 10
    EXCLUSION = 16
    NEGATION = 11
    SUBTRACT = 14
    DIVIDE = 15
    XOR = 13


__normalized = {
    BlendMode.NORMAL: True,
    BlendMode.MULTIPLY: True,
    BlendMode.DARKEN: True,
    BlendMode.LIGHTEN: True,
    BlendMode.ADD: False,
    BlendMode.COLOR_BURN: False,
    BlendMode.COLOR_DODGE: False,
    BlendMode.REFLECT: False,
    BlendMode.GLOW: False,
    BlendMode.OVERLAY: True,
    BlendMode.DIFFERENCE: True,
    BlendMode.NEGATION: True,
    BlendMode.SCREEN: True,
    BlendMode.XOR: True,
    BlendMode.SUBTRACT: False,
    BlendMode.DIVIDE: False,
    BlendMode.EXCLUSION: True,
    BlendMode.SOFT_LIGHT: True,
    BlendMode.HARD_LIGHT: True,
    BlendMode.VIVID_LIGHT: False,
    BlendMode.LINEAR_LIGHT: False,
    BlendMode.PIN_LIGHT: True,
    BlendMode.LINEAR_BURN: False,
}


def blend_mode_normalized(blend_mode: BlendMode) -> bool:
    """
    Returns whether the given blend mode is guaranteed to produce normalized results (value between 0 and 1).
    """
    return __normalized.get(blend_mode, False)


class ImageBlender:
    """Class for compositing images using different blending modes."""

    def __init__(self):
        self.modes = {
            BlendMode.NORMAL: self.__normal,
            BlendMode.MULTIPLY: self.__multiply,
            BlendMode.DARKEN: self.__darken,
            BlendMode.LIGHTEN: self.__lighten,
            BlendMode.ADD: self.__add,
            BlendMode.COLOR_BURN: self.__color_burn,
            BlendMode.COLOR_DODGE: self.__color_dodge,
            BlendMode.REFLECT: self.__reflect,
            BlendMode.GLOW: self.__glow,
            BlendMode.OVERLAY: self.__overlay,
            BlendMode.DIFFERENCE: self.__difference,
            BlendMode.NEGATION: self.__negation,
            BlendMode.SCREEN: self.__screen,
            BlendMode.XOR: self.__xor,
            BlendMode.SUBTRACT: self.__subtract,
            BlendMode.DIVIDE: self.__divide,
            BlendMode.EXCLUSION: self.__exclusion,
            BlendMode.SOFT_LIGHT: self.__soft_light,
            BlendMode.HARD_LIGHT: self.__hard_light,
            BlendMode.VIVID_LIGHT: self.__vivid_light,
            BlendMode.LINEAR_LIGHT: self.__linear_light,
            BlendMode.PIN_LIGHT: self.__pin_light,
            BlendMode.LINEAR_BURN: self.__linear_burn,
        }

    def apply_blend(
        self, a: np.ndarray, b: np.ndarray, blend_mode: BlendMode
    ) -> np.ndarray:
        return self.modes[blend_mode](a, b)

    def __normal(self, a: np.ndarray, _: np.ndarray) -> np.ndarray:
        return a

    def __multiply(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return a * b

    def __darken(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return np.minimum(a, b)

    def __lighten(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return np.maximum(a, b)

    def __add(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return a + b

    def __color_burn(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return np.where(
            a == 0, 0, np.maximum(0, (1 - ((1 - b) / np.maximum(0.0001, a))))
        )

    def __color_dodge(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return np.where(a == 1, 1, np.minimum(1, b / np.maximum(0.0001, (1 - a))))

    def __reflect(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return np.where(a == 1, 1, np.minimum(1, b * b / np.maximum(0.0001, 1 - a)))

    def __glow(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return np.where(b == 1, 1, np.minimum(1, a * a / np.maximum(0.0001, 1 - b)))

    def __overlay(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return np.where(b < 0.5, 2 * b * a, 1 - 2 * (1 - b) * (1 - a))

    def __difference(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return np.asarray(cv2.absdiff(a, b))

    def __negation(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return 1 - cv2.absdiff(1 - b, a)  # type: ignore

    def __screen(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return a + b - (a * b)  # type: ignore

    def __xor(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return normalize(
            np.bitwise_xor(to_uint8(a, normalized=True), to_uint8(b, normalized=True))
        )

    def __subtract(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return b - a

    def __divide(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return b / np.maximum(0.0001, a)

    def __exclusion(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return a * (1 - b) + b * (1 - a)

    def __soft_light(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        l = 2 * b * a + np.square(b) * (1 - 2 * a)
        h = np.sqrt(b) * (2 * a - 1) + 2 * b * (1 - a)
        return np.where(a <= 0.5, l, h)

    def __hard_light(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return np.where(a <= 0.5, 2 * a * b, 1 - 2 * (1 - a) * (1 - b))

    def __vivid_light(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return np.where(a <= 0.5, self.__color_burn(a, b), self.__color_dodge(a, b))

    def __linear_light(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return b + 2 * a - 1

    def __pin_light(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        x = 2 * a
        y = x - 1
        return np.where(b < y, y, np.where(b > x, x, b))

    def __linear_burn(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return a + b - 1


def blend_images(overlay: np.ndarray, base: np.ndarray, blend_mode: BlendMode):
    """
    Changes the given image to the background overlayed with the image.

    The 2 given images must be the same size and their values must be between 0 and 1.

    The returned image is guaranteed to have values between 0 and 1.

    If the 2 given images have a different number of channels, then the returned image
    will have maximum of the two.

    Only grayscale, RGB, and RGBA images are supported.
    """
    o_shape = get_h_w_c(overlay)
    b_shape = get_h_w_c(base)

    assert (
        o_shape[:2] == b_shape[:2]
    ), "The overlay and the base image must have the same size"

    def assert_sane(c: int, name: str):
        sane = c in (1, 3, 4)
        assert sane, f"The {name} has to be a grayscale, RGB, or RGBA image"

    o_channels = o_shape[2]
    b_channels = b_shape[2]

    assert_sane(o_channels, "overlay layer")
    assert_sane(b_channels, "base layer")

    blender = ImageBlender()
    target_c = max(o_channels, b_channels)
    needs_clipping = not blend_mode_normalized(blend_mode)

    if target_c == 4 and b_channels < 4:
        base = as_target_channels(base, 3)

        # The general algorithm below can be optimized because we know that b_a is 1
        o_a = np.dstack((overlay[:, :, 3],) * 3)
        o_rgb = overlay[:, :, :3]

        blend_rgb = blender.apply_blend(o_rgb, base, blend_mode)
        final_rgb = o_a * blend_rgb + (1 - o_a) * base  # type: ignore
        if needs_clipping:
            final_rgb = np.clip(final_rgb, 0, 1)

        return as_target_channels(final_rgb, 4)

    overlay = as_target_channels(overlay, target_c)
    base = as_target_channels(base, target_c)

    if target_c in (1, 3):
        # We don't need to do any alpha blending, so the images can blended directly
        result = blender.apply_blend(overlay, base, blend_mode)
        if needs_clipping:
            result = np.clip(result, 0, 1)
        return result

    # do the alpha blending for RGBA
    o_a = overlay[:, :, 3]
    b_a = base[:, :, 3]
    o_rgb = overlay[:, :, :3]
    b_rgb = base[:, :, :3]

    final_a = 1 - (1 - o_a) * (1 - b_a)

    blend_strength = o_a * b_a
    o_strength = o_a - blend_strength  # type: ignore
    b_strength = b_a - blend_strength  # type: ignore

    blend_rgb = blender.apply_blend(o_rgb, b_rgb, blend_mode)

    final_rgb = (
        (np.dstack((o_strength,) * 3) * o_rgb)
        + (np.dstack((b_strength,) * 3) * b_rgb)
        + (np.dstack((blend_strength,) * 3) * blend_rgb)
    )
    final_rgb /= np.maximum(np.dstack((final_a,) * 3), 0.0001)  # type: ignore
    final_rgb = np.clip(final_rgb, 0, 1)

    result = np.concatenate([final_rgb, np.expand_dims(final_a, axis=2)], axis=2)
    if needs_clipping:
        result = np.clip(result, 0, 1)
    return result

nodes/impl/blend.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import os
import sys
from enum import Enum
from typing import Tuple

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont

from ..utils.utils import get_h_w_c
from .image_utils import normalize, to_uint8


class CaptionPosition(Enum):
    BOTTOM = "bottom"
    TOP = "top"


def get_font_size(font: ImageFont.FreeTypeFont, text: str) -> Tuple[int, int]:
    """Get font [width, height] of the given text"""
    # (left, top, right, bottom)
    caption_bb = font.getbbox(text)
    font_width = caption_bb[2] - caption_bb[0]
    font_height = caption_bb[3] - caption_bb[1]
    return font_width, font_height


def add_caption(
    img: np.ndarray, caption: str, size: int, position: CaptionPosition
) -> np.ndarray:
    """Add caption with PIL"""
    fontsize = round(size * 0.8)
    if position is CaptionPosition.BOTTOM:
        img = cv2.copyMakeBorder(
            img, 0, size, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0, 1)
        )
    elif position is CaptionPosition.TOP:
        img = cv2.copyMakeBorder(
            img, size, 0, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0, 1)
        )
    else:
        raise RuntimeError(f"Unknown position {position}")

    pimg = Image.fromarray(to_uint8(img))
    font_path = os.path.join(
        os.path.dirname(sys.modules["__main__"].__file__), "fonts/Roboto-Light.ttf"  # type: ignore
    )
    font = ImageFont.truetype(font_path, fontsize)
    h, w, c = get_h_w_c(img)
    text_x = w // 2
    if position is CaptionPosition.BOTTOM:
        text_y = h - round(size / 2)
    elif position is CaptionPosition.TOP:
        text_y = round(size / 2)
    font_color = (255,) * c

    fw, _ = get_font_size(font, caption)
    # scale font size to fit image
    if fw > w:
        font = ImageFont.truetype(font_path, round(fontsize * w / fw))

    d = ImageDraw.Draw(pimg)
    d.text(
        (text_x, text_y),
        caption,
        font=font,
        anchor="mm",
        align="center",
        fill=font_color,
    )

    img = normalize(np.array(pimg))

    return img

nodes/impl/caption.py
vladmandic/sd-extension-chainner
import cv2
import numpy as np

from ..utils.utils import get_h_w_c
from .image_utils import as_2d_grayscale


def _luminance(img: np.ndarray) -> np.ndarray:
    """Returns the luminance of an image."""
    _, _, c = get_h_w_c(img)
    if c == 1:
        return as_2d_grayscale(img)
    if c == 2:
        return img[..., 0]
    return np.dot(img[..., :3], [0.2126, 0.7152, 0.0722])


def create_cas_mask(img: np.ndarray, kernel, bias: float = 2) -> np.ndarray:
    """
    Uses contrast adaptive sharpening's method to create a mask to interpolate between the original
    and sharpened image.

    `kernel` is an element create by `cv2.getStructuringElement`. It determines the shape and size
    of each pixel's neighborhood.

    `bias` is used to bias the mask towards the more sharpening. A value of 1 means no bias, a
    value greater than 1 biases the mask towards the sharpened image, a value less than 1 biases
    the mask towards the original image.

    Reference:
    Lou Kramer, FidelityFX CAS, AMD Developer Day 2019, https://gpuopen.com/wp-content/uploads/2019/07/FidelityFX-CAS.pptx
    https://www.shadertoy.com/view/wtlSWB#
    """
    assert bias > 0, "Bias must be greater than or equal to 0."

    l = _luminance(img)
    min_l = cv2.erode(l, kernel)
    max_l = cv2.dilate(l, kernel)
    min_d = np.minimum(1.0 - max_l, min_l, out=min_l)  # type: ignore
    max_l += 1e-8  # type: ignore
    min_d /= max_l
    mask = min_d
    if bias != 1:
        mask = np.power(mask, 1 / bias, out=mask)
    return mask


def cas_mix(
    img: np.ndarray,
    sharpened: np.ndarray,
    kernel,
    bias: float = 2,
) -> np.ndarray:
    mask = create_cas_mask(img, kernel, bias)
    _, _, c = get_h_w_c(sharpened)
    if c > 1:
        mask = np.dstack((mask,) * c)
    return img * (1 - mask) + sharpened * mask  # type: ignore

nodes/impl/cas.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import json
from typing import Iterable, List, Literal, Tuple, TypedDict, Union, cast

import numpy as np

from nodes.utils.utils import get_h_w_c

FloatLike = Union[np.floating, float]


def _norm(n: FloatLike) -> float:
    return max(0, min(float(n), 1))


ColorJsonKind = Literal["grayscale", "rgb", "rgba"]


class ColorJson(TypedDict):
    kind: ColorJsonKind
    values: List[float]


class Color:
    def __init__(self, value: Tuple[float, ...]) -> None:
        assert len(value) >= 1
        self.value: Tuple[float, ...] = value

    @property
    def channels(self) -> int:
        return len(self.value)

    @staticmethod
    def gray(gray: FloatLike) -> "Color":
        return Color((_norm(gray),))

    @staticmethod
    def bgr(value: Iterable[FloatLike]) -> "Color":
        t = tuple(map(_norm, value))
        assert len(t) == 3
        return Color(t)

    @staticmethod
    def bgra(value: Iterable[FloatLike]) -> "Color":
        t = tuple(map(_norm, value))
        assert len(t) == 4
        return Color(t)

    @staticmethod
    def from_1x1_image(img: np.ndarray) -> "Color":
        h, w, c = get_h_w_c(img)
        assert h == w == 1

        if c == 1:
            return Color.gray(img.flat[0])
        elif c == 3:
            return Color.bgr(img.flat)
        elif c == 4:
            return Color.bgra(img.flat)
        else:
            raise AssertionError("Only grayscale, RGB, and RGBA colors are supported.")

    @staticmethod
    def from_json(color_json: ColorJson | str) -> "Color":
        if isinstance(color_json, str):
            color_json = cast(ColorJson, json.loads(color_json))
        kind = color_json["kind"]
        values = color_json["values"]

        if kind == "grayscale":
            assert len(values) == 1
            return Color.gray(values[0])
        elif kind == "rgb":
            assert len(values) == 3
            return Color.bgr([values[2], values[1], values[0]])
        elif kind == "rgba":
            assert len(values) == 4
            return Color.bgra([values[2], values[1], values[0], values[3]])
        else:
            raise AssertionError(f"Unknown color kind {kind}")

    def to_1x1_image(self) -> np.ndarray:
        return self.to_image(1, 1)

    def to_image(self, width: int, height: int) -> np.ndarray:
        v = self.value
        if len(v) == 1:
            return np.full((height, width), v[0], dtype=np.float32)
        else:
            return np.full((height, width, len(v)), v, dtype=np.float32)

    def to_json(self) -> ColorJson:
        values = list(self.value)
        kind: ColorJsonKind
        if len(values) == 1:
            kind = "grayscale"
        elif len(values) == 3:
            kind = "rgb"
            values = [values[2], values[1], values[0]]
        elif len(values) == 4:
            kind = "rgba"
            values = [values[2], values[1], values[0], values[3]]
        else:
            raise AssertionError(f"Colors with {len(values)} are not supported.")

        return {"kind": kind, "values": values}

nodes/impl/color/color.py
vladmandic/sd-extension-chainner
from dataclasses import dataclass, field
from typing import Callable, Dict, Generic, Iterable, List, Set, Tuple, TypeVar, Union

import numpy as np
from nodes.log import logger

from .convert_data import color_spaces, color_spaces_or_detectors, conversions
from .convert_model import (
    ColorSpace,
    ColorSpaceDetector,
    Conversion,
    assert_input_channels,
    assert_output_channels,
)


def color_space_from_id(id_: int) -> ColorSpace:
    for c in color_spaces:
        if c.id == id_:
            return c
    raise ValueError(f"There is no color space with the id {id_}.")


def color_space_or_detector_from_id(id_: int) -> Union[ColorSpace, ColorSpaceDetector]:
    for c in color_spaces_or_detectors:
        if c.id == id_:
            return c
    raise ValueError(f"There is no color space with the id {id_}.")


T = TypeVar("T")


@dataclass(order=True)
class __ProcessingItem(Generic[T]):
    cost: int
    path: List[T] = field(compare=False)


def get_shortest_path(
    start: T,
    is_destination: Callable[[T], bool],
    get_next: Callable[[T], Iterable[Tuple[int, T]]],
) -> Union[List[T], None]:
    """A simple implementation of Dijkstra's"""

    processed: Set[T] = set()
    front: Dict[T, __ProcessingItem] = {
        start: __ProcessingItem(cost=0, path=[start]),
    }

    while len(front) > 0:
        best = None
        for x in front.values():
            if best is None:
                best = x
            elif x.cost < best.cost:
                best = x
        assert best is not None

        current = best.path[-1]
        del front[current]
        processed.add(current)

        if is_destination(current):
            return best.path

        for cost, to in get_next(current):
            cost = best.cost + cost
            old = front.get(to, None)
            if old is None:
                if to not in processed:
                    new_path = best.path.copy()
                    new_path.append(to)
                    front[to] = __ProcessingItem(cost=cost, path=new_path)
            else:
                if old.cost > cost:
                    old.cost = cost
                    old.path.clear()
                    old.path.extend(best.path)
                    old.path.append(to)


__conversions_map: Dict[ColorSpace, List[Conversion]] = {}
for conversion in conversions:
    l = __conversions_map.get(conversion.input, [])
    if len(l) == 0:
        __conversions_map[conversion.input] = l
    l.append(conversion)


def convert(
    img: np.ndarray,
    input_: Union[ColorSpace, ColorSpaceDetector],
    output: ColorSpace,
) -> np.ndarray:
    if isinstance(input_, ColorSpaceDetector):
        input_ = input_.detect(img)

    assert_input_channels(img, input_, output)

    if input_ == output:
        return img

    path = get_shortest_path(
        input_,
        is_destination=lambda i: i == output,
        get_next=lambda i: [(c.cost, c.output) for c in __conversions_map.get(i, [])],
    )

    if path is None:
        raise ValueError(f"Conversion {input_.name} -> {output.name} is not possible.")

    logger.debug(
        f"Converting color using the path {' -> '.join((x.name for x in path))}"
    )

    for i in range(1, len(path)):
        curr_in = path[i - 1]
        curr_out = path[i]

        conv = None
        for c in __conversions_map.get(curr_in, []):
            if c.output == curr_out:
                conv = c
                break
        assert conv is not None

        img = conv.convert(img)

    assert_output_channels(img, input_, output)
    return img

nodes/impl/color/convert.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import math
from typing import Dict, List

import cv2
import numpy as np

from ...utils.utils import get_h_w_c
from .convert_model import ColorSpace, ColorSpaceDetector, Conversion

GRAY = ColorSpace(0, "Gray", 1)
RGB = ColorSpace(1, "RGB", 3)
RGBA = ColorSpace(2, "RGBA", 4)
YUV = ColorSpace(3, "YUV", 3)
HSV = ColorSpace(4, "HSV", 3)
HSL = ColorSpace(5, "HSL", 3)
CMYK = ColorSpace(6, "CMYK", 4)
YUVA = ColorSpace(7, "YUVA", 4)
HSVA = ColorSpace(8, "HSVA", 4)
HSLA = ColorSpace(9, "HSLA", 4)
LAB = ColorSpace(10, "L*a*b*", 3)
LABA = ColorSpace(11, "L*a*b*A", 4)
LCH = ColorSpace(12, "L*C*h°", 3)
LCHA = ColorSpace(13, "L*C*h°A", 4)

RGB_LIKE = ColorSpaceDetector(1000, "RGB", [GRAY, RGB, RGBA])
YUV_LIKE = ColorSpaceDetector(1001, "YUV", [YUV, YUVA])
HSV_LIKE = ColorSpaceDetector(1002, "HSV", [HSV, HSVA])
HSL_LIKE = ColorSpaceDetector(1003, "HSL", [HSL, HSLA])
LAB_LIKE = ColorSpaceDetector(1004, "L*a*b*", [LAB, LABA])
LCH_LIKE = ColorSpaceDetector(1005, "L*C*h°", [LCH, LCHA])

ALPHA_PAIRS: Dict[ColorSpace, ColorSpace] = {
    RGB: RGBA,
    YUV: YUVA,
    HSV: HSVA,
    HSL: HSLA,
    LAB: LABA,
    LCH: LCHA,
}


def is_alpha_partner(c: ColorSpace) -> bool:
    """
    Whether this color space is returned by `get_alpha_partner` for some input.
    """
    return c in ALPHA_PAIRS.values()


def get_alpha_partner(c: ColorSpace) -> ColorSpace | None:
    """
    If the given color does NOT have an alpha channel and there exists another color space that this equivalent to the given one but does have an alpha, then the color space with an alpha channel will be returned, `None` otherwise.

    E.g. RGB will return RGBA and RGBA will return None.
    """
    return ALPHA_PAIRS.get(c, None)


color_spaces: List[ColorSpace] = [
    RGB,
    RGBA,
    GRAY,
    YUV,
    YUVA,
    HSV,
    HSVA,
    HSL,
    HSLA,
    CMYK,
    LAB,
    LABA,
    LCH,
    LCHA,
]
color_spaces_or_detectors: List[ColorSpace | ColorSpaceDetector] = [
    RGB_LIKE,
    GRAY,
    YUV_LIKE,
    HSV_LIKE,
    HSL_LIKE,
    CMYK,
    LAB_LIKE,
    LCH_LIKE,
]


def __rev3(image: np.ndarray) -> np.ndarray:
    c = get_h_w_c(image)[2]
    assert c == 3, "Expected a 3-channel image"
    return np.stack([image[:, :, 2], image[:, :, 1], image[:, :, 0]], axis=2)


def __rgb_to_hsv(img: np.ndarray) -> np.ndarray:
    img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    img[:, :, 0] /= 360  # type: ignore
    return __rev3(img)


def __hsv_to_rgb(img: np.ndarray) -> np.ndarray:
    img = __rev3(img)
    img[:, :, 0] *= 360
    return cv2.cvtColor(img, cv2.COLOR_HSV2BGR)


def __rgb_to_hsl(img: np.ndarray) -> np.ndarray:
    img = cv2.cvtColor(img, cv2.COLOR_BGR2HLS)
    h = img[:, :, 0] / 360  # type: ignore
    l = img[:, :, 1]
    s = img[:, :, 2]
    return cv2.merge((l, s, h))


def __hsl_to_rgb(img: np.ndarray) -> np.ndarray:
    h = img[:, :, 2] * 360
    s = img[:, :, 1]
    l = img[:, :, 0]
    return cv2.cvtColor(cv2.merge((h, l, s)), cv2.COLOR_HLS2BGR)


def __hsv_to_hsl(img: np.ndarray) -> np.ndarray:
    # the S and HSV and HSL are different, only the H is the same
    h = img[:, :, 2]
    hls = cv2.cvtColor(__hsv_to_rgb(img), cv2.COLOR_BGR2HLS)
    l = hls[:, :, 1]
    s = hls[:, :, 2]
    return cv2.merge((l, s, h))


def __hsl_to_hsv(img: np.ndarray) -> np.ndarray:
    # the S and HSV and HSL are different, only the H is the same
    h = img[:, :, 2]
    hsv = cv2.cvtColor(__hsl_to_rgb(img), cv2.COLOR_BGR2HSV)
    s = hsv[:, :, 1]
    v = hsv[:, :, 2]
    return cv2.merge((v, s, h))


def __rgb_to_cmyk(img: np.ndarray) -> np.ndarray:
    b, g, r = img[:, :, 0], img[:, :, 1], img[:, :, 2]
    maximum = np.max(img, axis=2)
    soft_max = np.maximum(maximum, 0.001)
    c = 1 - r / soft_max
    m = 1 - g / soft_max
    y = 1 - b / soft_max
    k = 1 - maximum
    return cv2.merge((y, m, c, k))


def __cmyk_to_rgb(img: np.ndarray) -> np.ndarray:
    y, m, c, k = img[:, :, 0], img[:, :, 1], img[:, :, 2], img[:, :, 3]
    maximum = 1 - k
    r = (1 - c) * maximum
    g = (1 - m) * maximum
    b = (1 - y) * maximum
    return cv2.merge((b, g, r))


def __rgb_to_lab(img: np.ndarray) -> np.ndarray:
    # 0≤L≤100 , −127≤a≤127, −127≤b≤127
    img = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    l = img[:, :, 0] / 100  # type: ignore
    a = (img[:, :, 1] + 127) / 254  # type: ignore
    b = (img[:, :, 2] + 127) / 254  # type: ignore
    return cv2.merge((b, a, l))


def __lab_to_rgb(img: np.ndarray) -> np.ndarray:
    # 0≤L≤100 , −127≤a≤127, −127≤b≤127
    l = img[:, :, 2] * 100
    a = img[:, :, 1] * 254 - 127
    b = img[:, :, 0] * 254 - 127
    return cv2.cvtColor(cv2.merge((l, a, b)), cv2.COLOR_LAB2BGR)


def __lab_to_lch(img: np.ndarray) -> np.ndarray:
    l = img[:, :, 2]
    # a and b must be centered at 0
    a = img[:, :, 1] - 0.5
    b = img[:, :, 0] - 0.5

    c = np.hypot(a, b)
    h = np.arctan2(b, a)

    # normalize C and h to [0,1]
    #
    # This is quite simple for h, but not so much for C. The problem is that
    # the maximum value of C depends on h. This is because a*,b* in [-1,1] form
    # a square of possible value around the origin. C nd h are simple polar
    # coordinates where h is the angel and C is the distance from origin. Since
    # the corners of a square are farther away from the origin than the mid
    # point of the sides of the a square, the maximum value of C depends on the
    # angle. E.g. for h=0, the maximum C value is 0.5 and for h=45°, it's
    # sqrt(0.5).
    #
    # One strategy would be simple use the maximum value for all possible h
    # values (which is sqrt(0.5)), but this has the problem that it is now
    # possible to create C,h value pairs that create invalid a*,b* values.
    # Ideally, all possible values C,h values should map to valid a*,b* values.
    #
    # To solve this problem, we calculate the maximum C value for the current
    # h angle and use that to normalize C. `Cmax` for a,b in [-1,1] is defined
    # as follows: `Cmax = 1/max(abs(cos(h)), abs(sin(h)))`. Since, we use a,b
    # in [-0.5,0.5], we just have to divide that value by 2.
    c_max = 0.5 / np.maximum(np.abs(np.sin(h)), np.abs(np.cos(h)))
    c = c / c_max
    h = h / (math.pi * 2) + 0.5

    return cv2.merge((h, c, l))


def __lch_to_lab(img: np.ndarray) -> np.ndarray:
    l = img[:, :, 2]

    # undo the c and h [0,1] normalization
    h = (img[:, :, 0] - 0.5) * (math.pi * 2)
    sin_h = np.sin(h)
    cos_h = np.cos(h)
    c_max = 0.5 / np.maximum(np.abs(sin_h), np.abs(cos_h))
    c = img[:, :, 1] * c_max

    a = c * cos_h + 0.5
    b = c * sin_h + 0.5
    return cv2.merge((b, a, l))


# The conversion loses one channel of information (e.g. the alpha channel, or a color channel)
__CHANNEL_LOST = 1000
# The conversion loses hue/chroma information in certain edge cases
__CHROMA_LOST = 100


conversions: List[Conversion] = [
    # RGB and grayscale
    Conversion(
        direction=(RGB, GRAY),
        convert=lambda i: cv2.cvtColor(i, cv2.COLOR_BGR2GRAY),
        cost=__CHANNEL_LOST * 2,
    ),
    Conversion(
        direction=(GRAY, RGB),
        convert=lambda i: cv2.cvtColor(i, cv2.COLOR_GRAY2BGR),
    ),
    Conversion(
        direction=(RGBA, GRAY),
        convert=lambda i: cv2.cvtColor(i, cv2.COLOR_BGRA2GRAY),
        cost=__CHANNEL_LOST * 3,
    ),
    Conversion(
        direction=(GRAY, RGBA),
        convert=lambda i: cv2.cvtColor(i, cv2.COLOR_GRAY2BGRA),
    ),
    # YUV
    Conversion(
        direction=(RGB, YUV),
        convert=lambda i: __rev3(cv2.cvtColor(i, cv2.COLOR_BGR2YUV)),
    ),
    Conversion(
        direction=(YUV, RGB),
        convert=lambda i: cv2.cvtColor(__rev3(i), cv2.COLOR_YUV2BGR),
        cost=__CHROMA_LOST,
    ),
    # HSV/HSL
    Conversion(
        direction=(RGB, HSV),
        convert=__rgb_to_hsv,
    ),
    Conversion(
        direction=(HSV, RGB),
        convert=__hsv_to_rgb,
        cost=__CHROMA_LOST,
    ),
    Conversion(
        direction=(RGB, HSL),
        convert=__rgb_to_hsl,
    ),
    Conversion(
        direction=(HSL, RGB),
        convert=__hsl_to_rgb,
        cost=__CHROMA_LOST,
    ),
    Conversion(
        direction=(HSV, HSL),
        convert=__hsv_to_hsl,
    ),
    Conversion(
        direction=(HSL, HSV),
        convert=__hsl_to_hsv,
    ),
    # CMYK
    Conversion(
        direction=(RGB, CMYK),
        convert=__rgb_to_cmyk,
    ),
    Conversion(
        direction=(CMYK, RGB),
        convert=__cmyk_to_rgb,
        cost=__CHROMA_LOST,
    ),
    # LAB
    Conversion(
        direction=(RGB, LAB),
        convert=__rgb_to_lab,
    ),
    Conversion(
        direction=(LAB, RGB),
        convert=__lab_to_rgb,
    ),
    # LCH
    Conversion(
        direction=(LAB, LCH),
        convert=__lab_to_lch,
    ),
    Conversion(
        direction=(LCH, LAB),
        convert=__lch_to_lab,
        cost=__CHROMA_LOST,
    ),
]


# Add conversions that can be generated because only alpha is different
for dir_3, dir_4 in ALPHA_PAIRS.items():
    assert dir_3.channels == 3
    assert dir_4.channels == 4

    # Add and remove the alpha channel
    conversions.append(
        Conversion(
            direction=(dir_3, dir_4),
            convert=lambda i: cv2.cvtColor(i, cv2.COLOR_BGR2BGRA),
        )
    )
    conversions.append(
        Conversion(
            direction=(dir_4, dir_3),
            convert=lambda i: cv2.cvtColor(i, cv2.COLOR_BGRA2BGR),
            cost=__CHANNEL_LOST,
        )
    )

__ALPHA_3_to_4 = dict(ALPHA_PAIRS)
for conv in list(conversions):
    in_4 = __ALPHA_3_to_4.get(conv.input)
    out_4 = __ALPHA_3_to_4.get(conv.output)

    if in_4 is not None and out_4 is not None:
        # if we have a conversion X -> Y, then we can generate XA -> YA
        # e.g. RGBA -> HSVA can be generated from RGB -> HSV

        def create_convert(old_conv: Conversion):
            def convert(img: np.ndarray) -> np.ndarray:
                color = img[:, :, :3]
                alpha = img[:, :, 3]
                return np.dstack((old_conv.convert(color), alpha))

            return convert

        conversions.append(
            Conversion(
                direction=(in_4, out_4),
                convert=create_convert(conv),
                cost=conv.cost,
            )
        )

nodes/impl/color/convert_data.py
vladmandic/sd-extension-chainner
from typing import Callable, Dict, Iterable, Tuple

import numpy as np

from ...utils.format import format_image_with_channels
from ...utils.utils import get_h_w_c


class ColorSpace:
    def __init__(self, id_: int, name: str, channels: int):
        assert 0 <= id_ and id_ < 256
        self.id = id_
        self.name = name
        self.channels = channels


class ColorSpaceDetector:
    def __init__(self, id_: int, name: str, color_spaces: Iterable[ColorSpace]) -> None:
        assert 1000 <= id_ and id_ < 2000
        self.id = id_
        self.name = name
        self.channel_map: Dict[int, ColorSpace] = {}
        for cs in color_spaces:
            assert cs.channels not in self.channel_map
            self.channel_map[cs.channels] = cs
        self.channels = list(self.channel_map.keys())

    def detect(self, image: np.ndarray) -> ColorSpace:
        c = get_h_w_c(image)[2]
        cs = self.channel_map.get(c, None)
        if cs is not None:
            return cs

        raise ValueError(
            f"Expected the input image for {self.name}"
            f" to be {format_image_with_channels(self.channels)}"
            f" but found {format_image_with_channels([c])}."
        )


def assert_input_channels(img: np.ndarray, input_: ColorSpace, output: ColorSpace):
    c = get_h_w_c(img)[2]
    if c != input_.channels:
        raise ValueError(
            f"Expected the input image for a {input_.name} -> {output.name} conversion"
            f" to be {format_image_with_channels([input_.channels])}"
            f" but found {format_image_with_channels([c])}."
        )


def assert_output_channels(result: np.ndarray, input_: ColorSpace, output: ColorSpace):
    c = get_h_w_c(result)[2]
    if c != output.channels:
        raise ValueError(
            f"Expected the output image for a {input_.name} -> {output.name} conversion"
            f" to be {format_image_with_channels([output.channels])}"
            f" but found {format_image_with_channels([c])}."
            f" This is an internal implementation error."
            f" Please report this as a bug."
        )


ConvertFn = Callable[[np.ndarray], np.ndarray]


class Conversion:
    def __init__(
        self,
        direction: Tuple[ColorSpace, ColorSpace],
        convert: ConvertFn,
        cost: int = 1,
    ):
        input_, output = direction
        assert input_ != output
        self.input: ColorSpace = input_
        self.output: ColorSpace = output
        self.__convert = convert
        assert cost >= 1
        self.cost: int = cost

    def convert(self, img: np.ndarray) -> np.ndarray:
        assert_input_channels(img, self.input, self.output)
        result = self.__convert(img)
        assert_output_channels(result, self.input, self.output)
        return result

nodes/impl/color/convert_model.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import numpy as np

__author__ = "Daniel Steinberg"
__copyright__ = "Copyright 2022, Daniel Steinberg"
__credits__ = ["Daniel Steinberg"]
__license__ = "MIT"
__version__ = "1.0.0"
__maintainer__ = "Daniel Steinberg"
__link__ = "https://github.com/dstein64/colortrans"


def matrix_sqrt(X):
    eig_val, eig_vec = np.linalg.eig(X)
    return eig_vec.dot(np.diag(np.sqrt(eig_val.clip(min=0)))).dot(eig_vec.T)


def linear_histogram_transfer(
    img: np.ndarray,
    ref_img: np.ndarray,
) -> np.ndarray:
    """
    Transfers the color distribution from the source to the target image
    using the Linear Histogram Matching..

    This implementation is based on to the Hertzmann, Aaron. "Algorithms
    for Rendering in Artistic Styles." Ph.D., New York University, 2001.
    """

    shape = img.shape

    # Convert HxWxC image to a (H*W)xC matrix.
    content = img.reshape(-1, shape[-1])
    reference = ref_img.reshape(-1, shape[-1])

    mu_content = np.mean(content, axis=0)
    mu_reference = np.mean(reference, axis=0)

    cov_content = np.cov(content, rowvar=False)
    cov_reference = np.cov(reference, rowvar=False)

    transfer = matrix_sqrt(cov_reference)
    sqrt_cov_content = matrix_sqrt(cov_content)

    if np.linalg.det(sqrt_cov_content) == 0:
        # Singular matrix: modify it by an arbitrary value before calculating the inverse matrix
        sqrt_cov_content += (
            np.identity(sqrt_cov_content.shape[-1], sqrt_cov_content.dtype) / 255.0
        )
    sqrt_cov_content_inv = np.linalg.inv(sqrt_cov_content)

    transfer = transfer.dot(sqrt_cov_content_inv)
    transfer = transfer.dot((content - mu_content).T).T
    transfer = transfer + mu_reference

    # Restore image dimensions.
    transfer = transfer.reshape(img.shape).clip(0, 1)

    return transfer

nodes/impl/color_transfer/linear_histogram.py
vladmandic/sd-extension-chainner
"""Fast color transfer method using mean and std"""

from __future__ import annotations

from enum import Enum

import cv2
import numpy as np

__author__ = "Adrian Rosebrock"
__copyright__ = "Copyright 2014, Adrian Rosebrock"
__credits__ = ["Adrian Rosebrock", "theflyingzamboni"]
__license__ = "MIT"
__version__ = "1.0.0"
__maintainer__ = "Adrian Rosebrock"
__link__ = "https://github.com/jrosebr1/color_transfer"


class TransferColorSpace(Enum):
    LAB = "L*a*b*"
    RGB = "RGB"


class OverflowMethod(Enum):
    CLIP = 1
    SCALE = 0


def image_stats(img: np.ndarray):
    """Get means and standard deviations of channels"""

    # Compute the mean and standard deviation of each channel
    channel_a, channel_b, channel_c = cv2.split(img)
    a_mean, a_std = (channel_a.mean(), channel_a.std())
    b_mean, b_std = (channel_b.mean(), channel_b.std())
    c_mean, c_std = (channel_c.mean(), channel_c.std())

    # Return the color statistics
    return a_mean, a_std, b_mean, b_std, c_mean, c_std


def min_max_scale(img: np.ndarray, new_range=(0, 255)):
    """Perform min-max scaling to a NumPy array"""

    # Get arrays current min and max
    mn = img.min()
    mx = img.max()

    # Check if scaling needs to be done to be in new_range
    if mn < new_range[0] or mx > new_range[1]:
        # Perform min-max scaling
        range_diff = new_range[1] - new_range[0]
        scaled = range_diff * (img - mn) / (mx - mn) + new_range[0]
    else:
        # Return array if already in range
        scaled = img

    return scaled


def scale_array(
    arr: np.ndarray,
    overflow_method: OverflowMethod,
    clip_min: int = 0,
    clip_max: int = 255,
) -> np.ndarray:
    """
    Trim NumPy array values to be in [0, 255] range with option of
    clipping or scaling.
    """

    if overflow_method == OverflowMethod.CLIP:
        scaled = np.clip(arr, clip_min, clip_max)
    else:
        scale_range = (max([arr.min(), clip_min]), min([arr.max(), clip_max]))
        scaled = min_max_scale(arr, new_range=scale_range)

    return scaled


def mean_std_transfer(
    img: np.ndarray,
    ref_img: np.ndarray,
    colorspace: TransferColorSpace,
    overflow_method: OverflowMethod,
    reciprocal_scale: bool = True,
) -> np.ndarray:
    """
    Transfers the color distribution from the source to the target image.
    Uses the mean and standard deviations of the specified
    colorspace. This implementation is (loosely) based on to the
    "Color Transfer between Images" paper by Reinhard et al., 2001.
    """

    a_clip_min, a_clip_max, b_clip_min, b_clip_max, c_clip_min, c_clip_max = (
        0,
        0,
        0,
        0,
        0,
        0,
    )

    # Convert the images from the RGB to L*a*b* color space, if necessary
    if colorspace == TransferColorSpace.LAB:
        a_clip_min, a_clip_max = (0, 100)
        b_clip_min, b_clip_max = (-127, 127)
        c_clip_min, c_clip_max = (-127, 127)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        ref_img = cv2.cvtColor(ref_img, cv2.COLOR_BGR2LAB)
    elif colorspace == TransferColorSpace.RGB:
        a_clip_min, a_clip_max = (0, 1)
        b_clip_min, b_clip_max = (0, 1)
        c_clip_min, c_clip_max = (0, 1)
        img = img[:, :, :3]
        ref_img = ref_img[:, :, :3]
    else:
        raise ValueError(f"Invalid color space {colorspace}")

    # Compute color statistics for the source and target images
    (
        a_mean_tar,
        a_std_tar,
        b_mean_tar,
        b_std_tar,
        c_mean_tar,
        c_std_tar,
    ) = image_stats(img)
    (
        a_mean_src,
        a_std_src,
        b_mean_src,
        b_std_src,
        c_mean_src,
        c_std_src,
    ) = image_stats(ref_img)

    # Subtract the means from the target image
    channel_a, channel_b, channel_c = cv2.split(img)
    channel_a -= a_mean_tar
    channel_b -= b_mean_tar
    channel_c -= c_mean_tar

    if reciprocal_scale:
        # Scale by the standard deviations using reciprocal of paper proposed factor
        channel_a = (a_std_src / a_std_tar) * channel_a
        channel_b = (b_std_src / b_std_tar) * channel_b
        channel_c = (c_std_src / c_std_tar) * channel_c
    else:
        # Scale by the standard deviations using paper proposed factor
        channel_a = (a_std_tar / a_std_src) * channel_a
        channel_b = (b_std_tar / b_std_src) * channel_b
        channel_c = (c_std_tar / c_std_src) * channel_c

    # Add in the source mean
    channel_a += a_mean_src
    channel_b += b_mean_src
    channel_c += c_mean_src

    # Clip/scale the pixel intensities to [clip_min, clip_max] if they fall
    # outside this range
    channel_a = scale_array(channel_a, overflow_method, a_clip_min, a_clip_max)
    channel_b = scale_array(channel_b, overflow_method, b_clip_min, b_clip_max)
    channel_c = scale_array(channel_c, overflow_method, c_clip_min, c_clip_max)

    # Merge the channels together, then convert back to the RGB color
    # space if necessary
    transfer = cv2.merge([channel_a, channel_b, channel_c])
    if colorspace == TransferColorSpace.LAB:
        transfer = cv2.cvtColor(transfer, cv2.COLOR_LAB2BGR)

    # Return the color transferred image
    return transfer

nodes/impl/color_transfer/mean_std.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import numpy as np

__author__ = "Daniel Steinberg"
__copyright__ = "Copyright 2022, Daniel Steinberg"
__credits__ = ["Daniel Steinberg"]
__license__ = "MIT"
__version__ = "1.0.0"
__maintainer__ = "Daniel Steinberg"
__link__ = "https://github.com/dstein64/colortrans"


def principal_color_transfer(
    img: np.ndarray,
    ref_img: np.ndarray,
) -> np.ndarray:
    """
    Transfers the color distribution from the source to the target image using
    Principal Component Color Matching.

    This implementation is based on:
    - Kotera, Hiroaki, Hung-Shing Chen, and Tetsuro Morimoto.
        "Object-to-Object Color Mapping by Image Segmentation." In Color Imaging:
        Device-Independent Color, Color Hardcopy, and Graphic Arts IV, 3648:148-57.
        SPIE, 1998.
    - Kotera, Hiroaki. "A Scene-Referred Color Transfer for Pleasant Imaging
        on Display." In IEEE International Conference on Image Processing 2005,
        2:II-5, 2005.
    """

    shape = img.shape

    # Convert HxWxC image to a (H*W)xC matrix.
    content = img.reshape(-1, shape[-1])
    reference = ref_img.reshape(-1, shape[-1])

    mu_content = np.mean(content, axis=0)
    mu_reference = np.mean(reference, axis=0)

    cov_content = np.cov(content, rowvar=False)
    cov_reference = np.cov(reference, rowvar=False)

    eigval_content, eigvec_content = np.linalg.eig(cov_content)
    eigval_reference, eigvec_reference = np.linalg.eig(cov_reference)

    # Division by 0 is forbidden: change null values to arbitrary values
    eigval_content = np.where(eigval_content == 0, 1e-42, eigval_content)
    eigval_factor = eigval_reference / eigval_content

    scaling = np.diag(np.sqrt(eigval_factor.clip(min=0)))
    transform = eigvec_reference.dot(scaling).dot(eigvec_content.T)
    transfer = (content - mu_content).dot(transform.T) + mu_reference

    # Restore image dimensions.
    transfer = transfer.reshape(shape).clip(0, 1)

    return transfer

nodes/impl/color_transfer/principal_color.py
vladmandic/sd-extension-chainner
from typing import Dict, Literal, Set, Union, cast

DxgiFormat = Literal[
    "UNKNOWN",
    "R32G32B32A32_TYPELESS",
    "R32G32B32A32_FLOAT",
    "R32G32B32A32_UINT",
    "R32G32B32A32_SINT",
    "R32G32B32_TYPELESS",
    "R32G32B32_FLOAT",
    "R32G32B32_UINT",
    "R32G32B32_SINT",
    "R16G16B16A16_TYPELESS",
    "R16G16B16A16_FLOAT",
    "R16G16B16A16_UNORM",
    "R16G16B16A16_UINT",
    "R16G16B16A16_SNORM",
    "R16G16B16A16_SINT",
    "R32G32_TYPELESS",
    "R32G32_FLOAT",
    "R32G32_UINT",
    "R32G32_SINT",
    "R32G8X24_TYPELESS",
    "D32_FLOAT_S8X24_UINT",
    "R32_FLOAT_X8X24_TYPELESS",
    "X32_TYPELESS_G8X24_UINT",
    "R10G10B10A2_TYPELESS",
    "R10G10B10A2_UNORM",
    "R10G10B10A2_UINT",
    "R11G11B10_FLOAT",
    "R8G8B8A8_TYPELESS",
    "R8G8B8A8_UNORM",
    "R8G8B8A8_UNORM_SRGB",
    "R8G8B8A8_UINT",
    "R8G8B8A8_SNORM",
    "R8G8B8A8_SINT",
    "R16G16_TYPELESS",
    "R16G16_FLOAT",
    "R16G16_UNORM",
    "R16G16_UINT",
    "R16G16_SNORM",
    "R16G16_SINT",
    "R32_TYPELESS",
    "D32_FLOAT",
    "R32_FLOAT",
    "R32_UINT",
    "R32_SINT",
    "R24G8_TYPELESS",
    "D24_UNORM_S8_UINT",
    "R24_UNORM_X8_TYPELESS",
    "X24_TYPELESS_G8_UINT",
    "R8G8_TYPELESS",
    "R8G8_UNORM",
    "R8G8_UINT",
    "R8G8_SNORM",
    "R8G8_SINT",
    "R16_TYPELESS",
    "R16_FLOAT",
    "D16_UNORM",
    "R16_UNORM",
    "R16_UINT",
    "R16_SNORM",
    "R16_SINT",
    "R8_TYPELESS",
    "R8_UNORM",
    "R8_UINT",
    "R8_SNORM",
    "R8_SINT",
    "A8_UNORM",
    "R1_UNORM",
    "R9G9B9E5_SHAREDEXP",
    "R8G8_B8G8_UNORM",
    "G8R8_G8B8_UNORM",
    "BC1_TYPELESS",
    "BC1_UNORM",
    "BC1_UNORM_SRGB",
    "BC2_TYPELESS",
    "BC2_UNORM",
    "BC2_UNORM_SRGB",
    "BC3_TYPELESS",
    "BC3_UNORM",
    "BC3_UNORM_SRGB",
    "BC4_TYPELESS",
    "BC4_UNORM",
    "BC4_SNORM",
    "BC5_TYPELESS",
    "BC5_UNORM",
    "BC5_SNORM",
    "B5G6R5_UNORM",
    "B5G5R5A1_UNORM",
    "B8G8R8A8_UNORM",
    "B8G8R8X8_UNORM",
    "R10G10B10_XR_BIAS_A2_UNORM",
    "B8G8R8A8_TYPELESS",
    "B8G8R8A8_UNORM_SRGB",
    "B8G8R8X8_TYPELESS",
    "B8G8R8X8_UNORM_SRGB",
    "BC6H_TYPELESS",
    "BC6H_UF16",
    "BC6H_SF16",
    "BC7_TYPELESS",
    "BC7_UNORM",
    "BC7_UNORM_SRGB",
    "AYUV",
    "Y410",
    "Y416",
    "NV12",
    "P010",
    "P016",
    "420_OPAQUE",
    "YUY2",
    "Y210",
    "Y216",
    "NV11",
    "AI44",
    "IA44",
    "P8",
    "A8P8",
    "B4G4R4A4_UNORM",
    "P208",
    "V208",
    "V408",
    "SAMPLER_FEEDBACK_MIN_MIP_OPAQUE,",
    "SAMPLER_FEEDBACK_MIP_REGION_USED_OPAQUE,",
    "FORCE_UINT",
]
"""
Modern DX10+ formats.

https://learn.microsoft.com/en-us/windows/win32/api/dxgiformat/ne-dxgiformat-dxgi_format
"""

LegacyFormat = Literal[
    "DXT1",
    "DXT3",
    "DXT5",
]
"""
Legacy DX9 formats. Those are the FourCC of those formats.
"""

DDSFormat = Union[DxgiFormat, LegacyFormat]

LEGACY_TO_DXGI: Dict[LegacyFormat, DxgiFormat] = {
    "DXT1": "BC1_UNORM",
    "DXT3": "BC2_UNORM",
    "DXT5": "BC3_UNORM",
}

SRGB_FORMATS: Set[DxgiFormat] = {
    "BC1_UNORM_SRGB",
    "BC2_UNORM_SRGB",
    "BC3_UNORM_SRGB",
    "BC7_UNORM_SRGB",
    "R8G8B8A8_UNORM_SRGB",
    "B8G8R8A8_UNORM_SRGB",
    "B8G8R8X8_UNORM_SRGB",
}

WITH_ALPHA: Set[DDSFormat] = {
    "R32G32B32A32_TYPELESS",
    "R32G32B32A32_FLOAT",
    "R32G32B32A32_UINT",
    "R32G32B32A32_SINT",
    "R16G16B16A16_TYPELESS",
    "R16G16B16A16_FLOAT",
    "R16G16B16A16_UNORM",
    "R16G16B16A16_UINT",
    "R16G16B16A16_SNORM",
    "R16G16B16A16_SINT",
    "R10G10B10A2_TYPELESS",
    "R10G10B10A2_UNORM",
    "R10G10B10A2_UINT",
    "R8G8B8A8_TYPELESS",
    "R8G8B8A8_UNORM",
    "R8G8B8A8_UNORM_SRGB",
    "R8G8B8A8_UINT",
    "R8G8B8A8_SNORM",
    "R8G8B8A8_SINT",
    "A8_UNORM",
    "BC1_TYPELESS",
    "BC1_UNORM",
    "BC1_UNORM_SRGB",
    "BC2_TYPELESS",
    "BC2_UNORM",
    "BC2_UNORM_SRGB",
    "BC3_TYPELESS",
    "BC3_UNORM",
    "BC3_UNORM_SRGB",
    "B5G5R5A1_UNORM",
    "B8G8R8A8_UNORM",
    "R10G10B10_XR_BIAS_A2_UNORM",
    "B8G8R8A8_TYPELESS",
    "B8G8R8A8_UNORM_SRGB",
    "BC7_TYPELESS",
    "BC7_UNORM",
    "BC7_UNORM_SRGB",
    "AYUV",
    "B4G4R4A4_UNORM",
    # legacy
    "DXT1",
    "DXT3",
    "DXT5",
}

BC7_FORMATS: Set[DxgiFormat] = {
    "BC7_TYPELESS",
    "BC7_UNORM",
    "BC7_UNORM_SRGB",
}
BC123_FORMATS: Set[DDSFormat] = {
    "BC1_TYPELESS",
    "BC1_UNORM",
    "BC1_UNORM_SRGB",
    "BC2_TYPELESS",
    "BC2_UNORM",
    "BC2_UNORM_SRGB",
    "BC3_TYPELESS",
    "BC3_UNORM",
    "BC3_UNORM_SRGB",
    # legacy
    "DXT1",
    "DXT3",
    "DXT5",
}

PREFER_DX9: Set[DDSFormat] = {
    "R8G8B8A8_UNORM",
    "B8G8R8A8_UNORM",
    "B5G5R5A1_UNORM",
    "B5G6R5_UNORM",
    "B8G8R8X8_UNORM",
    "R8G8_UNORM",
    "R8_UNORM",
}


def to_dxgi(f: DDSFormat) -> DxgiFormat:
    if f in LEGACY_TO_DXGI:
        return LEGACY_TO_DXGI[cast(LegacyFormat, f)]
    return cast(DxgiFormat, f)

nodes/impl/dds/format.py
vladmandic/sd-extension-chainner
import os
import platform
import shutil
import subprocess
import sys
import uuid
from tempfile import mkdtemp
from typing import List

import numpy as np
from nodes.log import logger

from ...utils.utils import split_file_path
from ..image_utils import cv_save_image
from .format import SRGB_FORMATS, DxgiFormat

__TEXCONV_DIR = os.path.join(
    os.path.dirname(sys.modules["__main__"].__file__), "texconv"  # type: ignore
)
__TEXCONV_EXE = os.path.join(__TEXCONV_DIR, "texconv.exe")


def __decode(b: bytes) -> str:
    try:
        return b.decode(encoding="iso8859-1")
    except:
        try:
            return b.decode(encoding="utf-8")
        except:
            return str(b)


def __run_texconv(args: List[str], error_message: str):
    if platform.system() != "Windows":
        # texconv is only supported on Windows.
        raise ValueError(
            "Texconv is only supported on Windows."
            " Reading and writing DDS files is only partially supported on other systems."
        )

    result = subprocess.run(
        [__TEXCONV_EXE, "-nologo", *args],
        check=False,
        capture_output=True,
    )

    if result.returncode != 0:
        output = (__decode(result.stdout) + __decode(result.stderr)).replace("\r", "")
        logger.error(f"chaiNNer: failed to run texconv: texconv: {__TEXCONV_EXE} args: {args} exit code: {result.returncode} output: {output}")
        raise ValueError(f"{error_message}: Code {result.returncode}: {output}")


def dds_to_png_texconv(path: str) -> str:
    """
    Converts the given DDS file to PNG by creating a temporary PNG file.
    """
    prefix = uuid.uuid4().hex
    _, basename, _ = split_file_path(path)

    tempdir = mkdtemp(prefix="chaiNNer-")

    __run_texconv(
        [
            "-f",
            "rgba",
            "-srgb",
            "-ft",
            "png",
            "-px",
            prefix,
            "-o",
            tempdir,
            path,
        ],
        "Unable to convert DDS",
    )

    return os.path.join(tempdir, prefix + basename + ".png")


def save_as_dds(
    path: str,
    image: np.ndarray,
    dds_format: DxgiFormat,
    mipmap_levels: int = 0,
    uniform_weighting: bool = False,
    dithering: bool = False,
    minimal_compression: bool = False,
    maximum_compression: bool = False,
    dx9: bool = False,
    separate_alpha: bool = False,
):
    """
    Saves an image as DDS using texconv.
    See the following page for more information on save options:
    https://github.com/Microsoft/DirectXTex/wiki/Texconv
    """
    target_dir, name, ext = split_file_path(path)

    assert ext == ".dds", "The file to save must end with '.dds'"

    tempDir = mkdtemp(prefix="chaiNNer-")

    try:
        tempPng = os.path.join(tempDir, f"{name}.png")
        cv_save_image(tempPng, image, [])

        args = [
            "-y",
            "-f",
            dds_format,
            "-dx9" if dx9 else "-dx10",
            "-m",
            str(mipmap_levels),
            # use texconv to directly produce the target file
            "-o",
            target_dir,
        ]

        bc = ""
        bc += "u" if uniform_weighting else ""
        bc += "d" if dithering else ""
        bc += "q" if minimal_compression else ""
        bc += "x" if maximum_compression else ""
        if bc:
            args.extend(["-bc", f"-{bc}"])

        if dds_format in SRGB_FORMATS:
            args.append("-srgbi")

        if separate_alpha:
            args.append("-sepalpha")

        args.append(tempPng)
        __run_texconv(args, "Unable to write DDS")
    finally:
        shutil.rmtree(tempDir)

nodes/impl/dds/texconv.py
vladmandic/sd-extension-chainner
from enum import Enum
from typing import Dict

from chainner_ext import DiffusionAlgorithm


class ErrorDiffusionMap(Enum):
    FLOYD_STEINBERG = "FS"
    JARVIS_ET_AL = "JJN"
    STUCKI = "ST"
    ATKINSON = "A"
    BURKES = "B"
    SIERRA = "S"
    TWO_ROW_SIERRA = "S2"
    SIERRA_LITE = "SL"


ERROR_PROPAGATION_MAP_LABELS = {
    ErrorDiffusionMap.FLOYD_STEINBERG: "Floyd-Steinberg",
    ErrorDiffusionMap.JARVIS_ET_AL: "Jarvis, Judice, and Ninke",
    ErrorDiffusionMap.STUCKI: "Stucki",
    ErrorDiffusionMap.ATKINSON: "Atkinson",
    ErrorDiffusionMap.BURKES: "Burkes",
    ErrorDiffusionMap.SIERRA: "Sierra",
    ErrorDiffusionMap.TWO_ROW_SIERRA: "Two Row Sierra",
    ErrorDiffusionMap.SIERRA_LITE: "Sierra Lite",
}

DIFFUSION_ALGORITHM_MAP: Dict[ErrorDiffusionMap, DiffusionAlgorithm] = {
    ErrorDiffusionMap.FLOYD_STEINBERG: DiffusionAlgorithm.FloydSteinberg,
    ErrorDiffusionMap.JARVIS_ET_AL: DiffusionAlgorithm.JarvisJudiceNinke,
    ErrorDiffusionMap.STUCKI: DiffusionAlgorithm.Stucki,
    ErrorDiffusionMap.ATKINSON: DiffusionAlgorithm.Atkinson,
    ErrorDiffusionMap.BURKES: DiffusionAlgorithm.Burkes,
    ErrorDiffusionMap.SIERRA: DiffusionAlgorithm.Sierra,
    ErrorDiffusionMap.TWO_ROW_SIERRA: DiffusionAlgorithm.TwoRowSierra,
    ErrorDiffusionMap.SIERRA_LITE: DiffusionAlgorithm.SierraLite,
}


# https://en.wikipedia.org/wiki/Ordered_dithering


class ThresholdMap(Enum):
    BAYER_2 = "B2"
    BAYER_4 = "B4"
    BAYER_8 = "B8"
    BAYER_16 = "B16"


THRESHOLD_MAP_LABELS = {
    ThresholdMap.BAYER_2: "Bayer 2x2",
    ThresholdMap.BAYER_4: "Bayer 4x4",
    ThresholdMap.BAYER_8: "Bayer 8x8",
    ThresholdMap.BAYER_16: "Bayer 16x16",
}

nodes/impl/dithering/constants.py
vladmandic/sd-extension-chainner
import cv2
import numpy as np

from ..image_utils import MAX_VALUES_BY_DTYPE, as_3d


def _as_float32(image: np.ndarray) -> np.ndarray:
    if image.dtype == np.float32:
        return image
    max_value = MAX_VALUES_BY_DTYPE[image.dtype.name]
    return image.astype(np.float32) / max_value


def distinct_colors_palette(image: np.ndarray) -> np.ndarray:
    image = as_3d(image)
    return np.unique(image.reshape((-1, image.shape[2])), axis=0).reshape(
        (1, -1, image.shape[2])
    )


def kmeans_palette(image: np.ndarray, num_colors: int) -> np.ndarray:
    image = as_3d(image)
    flat_image = _as_float32(image.reshape((-1, image.shape[2])))

    max_iter = 10
    epsilon = 1.0
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, max_iter, epsilon)

    attempts = 10
    cv2.setRNGSeed(0)
    _, _, center = cv2.kmeans(
        flat_image, num_colors, None, criteria, attempts, cv2.KMEANS_PP_CENTERS  # type: ignore
    )

    return center.reshape((1, -1, image.shape[2]))


class MedianCutBucket:
    def __init__(self, data: np.ndarray):
        self.data = data
        self.n_pixels, self.n_channels = data.shape
        self.min_values = np.min(data, axis=0)
        self.max_values = np.max(data, axis=0)
        self.channel_ranges = self.max_values - self.min_values
        self.biggest_range = np.max(self.channel_ranges)

    def split(self):
        widest_channel = np.argmax(self.channel_ranges)
        median = np.median(self.data[:, widest_channel])
        mask = self.data[:, widest_channel] > median
        if mask.sum() == 0:
            mean = np.mean(self.data[:, widest_channel])
            mask = self.data[:, widest_channel] > mean
        return MedianCutBucket(self.data[mask is True]), MedianCutBucket(
            self.data[mask is False]
        )

    def average(self):
        return np.mean(self.data, axis=0)


def median_cut_palette(image: np.ndarray, num_colors: int) -> np.ndarray:
    image = as_3d(image)
    flat_image = _as_float32(image.reshape((-1, image.shape[2])))

    buckets = [MedianCutBucket(flat_image)]
    while len(buckets) < num_colors:
        bucket_idx, bucket = max(enumerate(buckets), key=lambda x: x[1].biggest_range)
        if bucket.biggest_range == 0:
            break
        buckets.pop(bucket_idx)

        buckets.extend(bucket.split())

    return np.stack([bucket.average() for bucket in buckets], axis=0).reshape(
        (1, -1, image.shape[2])
    )

nodes/impl/dithering/palette.py
vladmandic/sd-extension-chainner
import math

import numpy as np


def horizontal_gradient(img: np.ndarray):
    x = np.arange(img.shape[1])
    p = x / (img.shape[1] - 1)
    img[:, :] = p.reshape((1, -1))


def vertical_gradient(img: np.ndarray):
    x = np.arange(img.shape[0])
    p = x / (img.shape[0] - 1)
    img[:, :] = p.reshape((-1, 1))


def diagonal_gradient(img: np.ndarray, angle: float, width: float):
    center = np.array([img.shape[0], img.shape[1]], dtype=np.float32) / 2
    direction = np.array([np.cos(angle), np.sin(angle)], dtype=np.float32)

    start = center - direction * width / 2

    pixels = np.array(
        [[(r, c) for r in range(img.shape[0]) for c in range(img.shape[1])]]
    )
    projection = (pixels - start).dot(direction)
    p = np.clip((projection / width).ravel(), 0, 1)
    img[:] = p.reshape(img.shape)


def radial_gradient(
    img: np.ndarray, inner_radius_percent: float = 0, outer_radius_percent: float = 1
):
    inner_radius = inner_radius_percent * img.shape[1] / 2
    outer_radius = outer_radius_percent * img.shape[1] / 2

    center = np.array(img.shape[:2], dtype="float32") / 2
    pixels = np.array(
        [(r, c) for r in range(img.shape[0]) for c in range(img.shape[1])]
    )
    distance = np.sqrt(np.sum((pixels - center) ** 2, axis=1))
    p = (distance - inner_radius) / (outer_radius - inner_radius)
    img[:] = p.reshape(img.shape)


def conic_gradient(img: np.ndarray, rotation: float = 0):
    if rotation > np.pi:
        rotation -= 2 * np.pi
    if rotation < -np.pi:
        rotation += 2 * np.pi

    center = np.array(img.shape[:2], dtype="float32") / 2
    pixels = np.array(
        [(r, c) for r in range(img.shape[0]) for c in range(img.shape[1])]
    )
    angles = np.arctan2(pixels[:, 0] - center[0], pixels[:, 1] - center[1]) + rotation
    angles[angles < 0] += 2 * np.pi
    p = angles / math.pi / 2
    img[:] = p.reshape(img.shape)

nodes/impl/gradients.py
vladmandic/sd-extension-chainner
def get_opencv_formats():
    return [
        # Bitmaps
        ".bmp",
        ".dib",
        # JPEG
        ".jpg",
        ".jpeg",
        ".jpe",
        ".jp2",
        # PNG, WebP, Tiff
        ".png",
        ".webp",
        ".tif",
        ".tiff",
        # Portable image format
        ".pbm",
        ".pgm",
        ".ppm",
        ".pxm",
        ".pnm",
        # Sun Rasters
        ".sr",
        ".ras",
        # OpenEXR
        ".exr",
        # Radiance HDR
        ".hdr",
        ".pic",
    ]


def get_pil_formats():
    return [
        # Bitmaps
        ".bmp",
        ".dib",
        ".xbm",
        # DDS
        ".dds",
        # EPS
        ".eps",
        # GIF
        # ".gif",
        # Icons
        ".icns",
        ".ico",
        # JPEG
        ".jpg",
        ".jpeg",
        ".jfif",
        ".jp2",
        ".jpx",
        # Randoms
        ".msp",
        ".pcx",
        ".sgi",
        # PNG, WebP, TIFF
        ".png",
        ".webp",
        ".tiff",
        # APNG
        # ".apng",
        # Portable image format
        ".pbm",
        ".pgm",
        ".ppm",
        ".pnm",
        # TGA
        ".tga",
    ]


def get_available_image_formats():
    all_formats = [*get_opencv_formats(), *get_pil_formats()]
    return sorted(set(all_formats))

nodes/impl/image_formats.py
vladmandic/sd-extension-chainner
from typing import Callable

import numpy as np
from typing_extensions import Concatenate, ParamSpec

ImageOp = Callable[[np.ndarray], np.ndarray]
"""
An image processing operation that takes an image and produces a new image.

The given image is guaranteed to *not* be modified.
"""


def clipped(op: ImageOp) -> ImageOp:
    """
    Ensures that all values in the returned image are between 0 and 1.
    """
    return lambda i: np.clip(op(i), 0, 1)


P = ParamSpec("P")


def to_op(fn: Callable[Concatenate[np.ndarray, P], np.ndarray]) -> Callable[P, ImageOp]:
    """
    Applies a form of currying to convert the given function into a constructor for an image operation.

    Example: Simple resize method could be defined as follows: `resize(np.ndarray, Size2D) -> np.ndarray`.
    It takes an image and its new size and returns the resized image.
    If we want to convert it to an image operation, we have to create a function with the following signature: `resize_op(Size2D) -> ImageOp`.
    The implementation of this function would be rather simple, it would simply take all arguments of `resize` except for the image like this:
    ```py
    def resize_op(size: Size2D) -> ImageOp:
        return lambda img: resize(img, size)
    ```
    `to_op` does exactly this transformation, but for any number of arguments.

    Note: This only works if the input image is the first argument of the given function.
    """

    def p(*args: P.args, **kwargs: P.kwargs) -> ImageOp:
        return lambda i: fn(i, *args, **kwargs)

    return p

nodes/impl/image_op.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import itertools
import math
import os
import random
import string
from enum import Enum
from typing import List, Tuple

import cv2
import numpy as np

from ..utils.utils import Padding, get_h_w_c, split_file_path
from .color.color import Color

MAX_VALUES_BY_DTYPE = {
    np.dtype("int8").name: 127,
    np.dtype("uint8").name: 255,
    np.dtype("int16").name: 32767,
    np.dtype("uint16").name: 65535,
    np.dtype("int32").name: 2147483647,
    np.dtype("uint32").name: 4294967295,
    np.dtype("int64").name: 9223372036854775807,
    np.dtype("uint64").name: 18446744073709551615,
    np.dtype("float32").name: 1.0,
    np.dtype("float64").name: 1.0,
}


class FillColor(Enum):
    AUTO = -1
    BLACK = 0
    TRANSPARENT = 1

    def get_color(self, channels: int):
        """Select how to fill negative space that results from rotation"""

        if self == FillColor.AUTO:
            fill_color = (0,) * channels
        elif self == FillColor.BLACK:
            fill_color = (0,) * channels if channels < 4 else (0, 0, 0, 1)
        else:
            fill_color = (0, 0, 0, 0)

        return fill_color


class FlipAxis(Enum):
    HORIZONTAL = 1
    VERTICAL = 0
    BOTH = -1
    NONE = 2

    def flip(self, img: np.ndarray) -> np.ndarray:
        if self == FlipAxis.NONE:
            return img
        return cv2.flip(img, self.value)


class BorderType(Enum):
    REFLECT_MIRROR = 4
    WRAP = 3
    REPLICATE = 1
    BLACK = 0
    WHITE = 6
    TRANSPARENT = 5
    CUSTOM_COLOR = 7


class NormalMapType(Enum):
    DIRECTX = "DirectX"
    OPENGL = "OpenGL"
    OCTAHEDRAL = "Octahedral"


def convert_to_BGRA(img: np.ndarray, in_c: int) -> np.ndarray:
    assert in_c in (1, 3, 4), f"Number of channels ({in_c}) unexpected"
    if in_c == 1:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGRA)
    elif in_c == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)

    return img.copy()


def _get_iinfo(img: np.ndarray) -> np.iinfo | None:
    try:
        return np.iinfo(img.dtype)
    except:
        return None


def normalize(img: np.ndarray) -> np.ndarray:
    if img.dtype != np.float32:
        info = _get_iinfo(img)
        img = img.astype(np.float32)

        if info is not None:
            img /= info.max
            if info.min == 0:
                # we don't need to clip
                return img

        # we own `img`, so it's okay to write to it
        return np.clip(img, 0, 1, out=img)

    return np.clip(img, 0, 1)


def to_uint8(
    img: np.ndarray,
    normalized=False,
) -> np.ndarray:
    """
    Returns a new uint8 image with the given image data.

    If `normalized` is `False`, then the image will be normalized before being converted to uint8.
    """
    if img.dtype == np.uint8:
        return img.copy()

    if not normalized or img.dtype != np.float32:
        img = normalize(img)

    return (img * 255).round().astype(np.uint8)


def to_uint16(
    img: np.ndarray,
    normalized=False,
) -> np.ndarray:
    """
    Returns a new uint8 image with the given image data.

    If `normalized` is `False`, then the image will be normalized before being converted to uint8.
    """
    if img.dtype == np.uint16:
        return img.copy()

    if not normalized or img.dtype != np.float32:
        img = normalize(img)

    return (img * 65535).round().astype(np.uint16)


def shift(img: np.ndarray, amount_x: int, amount_y: int, fill: FillColor) -> np.ndarray:
    c = get_h_w_c(img)[2]
    if fill == FillColor.TRANSPARENT:
        img = convert_to_BGRA(img, c)
    fill_color = fill.get_color(c)

    h, w, _ = get_h_w_c(img)
    translation_matrix = np.asfarray(
        [[1, 0, amount_x], [0, 1, amount_y]], dtype=np.float32
    )
    img = cv2.warpAffine(
        img,
        translation_matrix,
        (w, h),
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=fill_color,
    )

    return img


def as_2d_grayscale(img: np.ndarray) -> np.ndarray:
    """Given a grayscale image, this returns an image with 2 dimensions (image.ndim == 2)."""
    if img.ndim == 2:
        return img
    if img.ndim == 3 and img.shape[2] == 1:
        return img[:, :, 0]
    raise AssertionError(f"Invalid image shape {img.shape}")


def as_3d(img: np.ndarray) -> np.ndarray:
    """Given a grayscale image, this returns an image with 3 dimensions (image.ndim == 3)."""
    if img.ndim == 2:
        return np.expand_dims(img.copy(), axis=2)
    return img


def as_target_channels(
    img: np.ndarray, target_c: int, narrowing: bool = False
) -> np.ndarray:
    """
    Given a number of target channels (either 1, 3, or 4), this convert the given image
    to an image with that many channels. If the given image already has the correct
    number of channels, it will be returned as is.

    Narrowing conversions are only supported if narrowing is True.
    """
    c = get_h_w_c(img)[2]

    if c == target_c == 1:
        return as_2d_grayscale(img)
    if c == target_c:
        return img

    if not narrowing:
        assert (
            c < target_c
        ), f"Narrowing is false, image channels ({c}) must be less than target channels ({target_c})"

    if c == 1:
        if target_c == 3:
            return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        if target_c == 4:
            return cv2.cvtColor(img, cv2.COLOR_GRAY2BGRA)

    if c == 3:
        if target_c == 1:
            return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        if target_c == 4:
            return cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)

    if c == 4:
        if target_c == 1:
            return cv2.cvtColor(img, cv2.COLOR_BGRA2GRAY)
        if target_c == 3:
            return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)

    raise ValueError(f"Unable to convert {c} channel image to {target_c} channel image")


def create_border(
    img: np.ndarray,
    border_type: BorderType,
    border: Padding,
    color: Color | None = None,
) -> np.ndarray:
    """
    Returns a new image with a specified border.
    """

    if border.empty:
        return img

    _, _, c = get_h_w_c(img)
    if c == 4 and border_type == BorderType.BLACK:
        value = (0.0, 0.0, 0.0, 1.0)
    else:
        value = (0.0,)

    cv_border_type: int = border_type.value
    if border_type == BorderType.TRANSPARENT:
        cv_border_type = cv2.BORDER_CONSTANT
        value = (0.0,)
        img = as_target_channels(img, 4)
    elif border_type == BorderType.WHITE:
        cv_border_type = cv2.BORDER_CONSTANT
        value = (1.0,) * c
    elif border_type == BorderType.CUSTOM_COLOR:
        assert (
            color is not None
        ), "Creating a border with a custom color requires supplying a custom color."

        # widen image or color to make them compatible
        if color.channels > c:
            img = as_target_channels(img, color.channels)
        elif c > color.channels:
            color = Color.from_1x1_image(as_target_channels(color.to_1x1_image(), c))

        cv_border_type = cv2.BORDER_CONSTANT
        value = color.value

    return cv2.copyMakeBorder(
        img,
        top=border.top,
        left=border.left,
        right=border.right,
        bottom=border.bottom,
        borderType=cv_border_type,
        value=value,
    )


def calculate_ssim(
    img1: np.ndarray,
    img2: np.ndarray,
) -> float:
    """Calculates mean localized Structural Similarity Index (SSIM)
    between two images."""

    C1 = 0.01**2
    C2 = 0.03**2

    kernel = cv2.getGaussianKernel(11, 1.5)
    window = np.outer(kernel, kernel.transpose())  # type: ignore

    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]
    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
    mu1_sq = np.power(mu1, 2)
    mu2_sq = np.power(mu2, 2)
    mu1_mu2 = np.multiply(mu1, mu2)
    sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2

    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / (
        (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
    )

    return float(np.mean(ssim_map))


def cv_save_image(path: str, img: np.ndarray, params: List[int]):
    """
    A light wrapper around `cv2.imwrite` to support non-ASCII paths.
    """

    # Write image with opencv if path is ascii, since imwrite doesn't support unicode
    # This saves us from having to keep the image buffer in memory, if possible
    if path.isascii():
        cv2.imwrite(path, img, params)
    else:
        dirname, _, extension = split_file_path(path)
        try:
            temp_filename = f'temp-{"".join(random.choices(string.ascii_letters, k=16))}.{extension}'
            full_temp_path = os.path.join(dirname, temp_filename)
            cv2.imwrite(full_temp_path, img, params)
            os.rename(full_temp_path, path)
        except:
            _, buf_img = cv2.imencode(f".{extension}", img, params)
            with open(path, "wb") as outf:
                outf.write(buf_img)  # type: ignore


def cartesian_product(arrays: List[np.ndarray]) -> np.ndarray:
    """
    Returns the cartesian product of the given arrays. Good for initializing coordinates, for example.

    This is cartesian_product_transpose_pp from this following SO post by Paul Panzer:
    https://stackoverflow.com/questions/11144513/cartesian-product-of-x-and-y-array-points-into-single-array-of-2d-points/49445693#49445693
    """
    #
    la = len(arrays)
    dtype = np.result_type(*arrays)
    arr = np.empty((la, *map(len, arrays)), dtype=dtype)
    idx = slice(None), *itertools.repeat(None, la)
    for i, a in enumerate(arrays):
        arr[i, ...] = a[idx[: la - i]]
    return arr.reshape(la, -1).T


def fast_gaussian_blur(
    img: np.ndarray,
    sigma_x: float,
    sigma_y: float | None = None,
) -> np.ndarray:
    """
    Computes a channel-wise gaussian blur of the given image using a fast approximation.

    The maximum error of the approximation is guaranteed to be less than 0.1%.
    In addition to that, the error is guaranteed to be smoothly distributed across the image.
    There are no sudden spikes in error anywhere.

    Specifically, the method is implemented by downsampling the image, blurring the downsampled
    image, and then upsampling the blurred image. This is much faster than blurring the full image.
    Unfortunately, OpenCV's `resize` method has unfortunate artifacts when upscaling, so we
    apply a small gaussian blur to the image after upscaling to smooth out the artifacts. This
    single step almost doubles the runtime of the method, but it is still much faster than
    blurring the full image.
    """
    if sigma_y is None:
        sigma_y = sigma_x
    if sigma_x == 0 or sigma_y == 0:
        return img.copy()

    h, w, _ = get_h_w_c(img)

    def get_scale_factor(sigma: float) -> float:
        if sigma < 11:
            return 1
        if sigma < 15:
            return 1.25
        if sigma < 20:
            return 1.5
        if sigma < 25:
            return 2
        if sigma < 30:
            return 2.5
        if sigma < 50:
            return 3
        if sigma < 100:
            return 4
        if sigma < 200:
            return 6
        return 8

    def get_sizing(size: int, sigma: float, f: float) -> Tuple[int, float, float]:
        """
        Return the size of the downsampled image, the sigma of the downsampled gaussian blur,
        and the sigma of the upscaled gaussian blur.
        """
        if f <= 1:
            # just use simple gaussian, the error is too large otherwise
            return size, 0, sigma

        size_down = math.ceil(size / f)
        f = size / size_down
        sigma_up = f
        sigma_down = math.sqrt(sigma**2 - sigma_up**2) / f
        return size_down, sigma_down, sigma_up

    # Handling different sigma values for x and y is difficult, so we take the easy way out
    # and just use the smaller one. There are potentially better ways of combining them, but
    # this is good enough for now.
    scale_factor = min(get_scale_factor(sigma_x), get_scale_factor(sigma_y))
    h_down, y_down_sigma, y_up_sigma = get_sizing(h, sigma_y, scale_factor)
    w_down, x_down_sigma, x_up_sigma = get_sizing(w, sigma_x, scale_factor)

    if h != h_down or w != w_down:
        # downsampled gaussian blur
        img = cv2.resize(img, (w_down, h_down), interpolation=cv2.INTER_AREA)
        img = cv2.GaussianBlur(
            img,
            (0, 0),
            sigmaX=x_down_sigma,
            sigmaY=y_down_sigma,
            borderType=cv2.BORDER_REFLECT,
        )
        img = cv2.resize(img, (w, h), interpolation=cv2.INTER_LINEAR)

    if x_up_sigma != 0 or y_up_sigma != 0:
        # post blur to smooth out artifacts
        img = cv2.GaussianBlur(
            img,
            (0, 0),
            sigmaX=x_up_sigma,
            sigmaY=y_up_sigma,
            borderType=cv2.BORDER_REFLECT,
        )

    return img

nodes/impl/image_utils.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import gc

import numpy as np

try:
    from ncnn_vulkan import ncnn

    use_gpu = True
except ImportError:
    from ncnn import ncnn

    use_gpu = False
from nodes.log import logger

from ...utils.utils import get_h_w_c
from ..image_utils import to_uint8
from ..upscale.auto_split import Split, Tiler, auto_split


def ncnn_auto_split(
    img: np.ndarray,
    net,
    input_name: str,
    output_name: str,
    blob_vkallocator,
    staging_vkallocator,
    tiler: Tiler,
) -> np.ndarray:
    def upscale(img: np.ndarray, _):
        ex = net.create_extractor()
        if use_gpu:
            ex.set_blob_vkallocator(blob_vkallocator)
            ex.set_workspace_vkallocator(blob_vkallocator)
            ex.set_staging_vkallocator(staging_vkallocator)
        # ex.set_light_mode(True)
        try:
            lr_c = get_h_w_c(img)[2]
            lr_img_fix = to_uint8(img)
            if lr_c == 1:
                pixel_type = ncnn.Mat.PixelType.PIXEL_GRAY
            elif lr_c == 3:
                pixel_type = ncnn.Mat.PixelType.PIXEL_RGB
            else:
                pixel_type = ncnn.Mat.PixelType.PIXEL_RGBA
            mat_in = ncnn.Mat.from_pixels(
                lr_img_fix,
                pixel_type,
                lr_img_fix.shape[1],
                lr_img_fix.shape[0],
            )
            mean_vals = []
            norm_vals = [1 / 255.0] * lr_c
            mat_in.substract_mean_normalize(mean_vals, norm_vals)
            ex.input(input_name, mat_in)
            _, mat_out = ex.extract(output_name)
            result = np.array(mat_out).transpose(1, 2, 0).astype(np.float32)
            del ex, mat_in, mat_out
            gc.collect()
            if use_gpu:
                # Clear VRAM
                blob_vkallocator.clear()
                staging_vkallocator.clear()
            return result
        except Exception as e:
            if "vkQueueSubmit" in str(e):
                ex = None
                del ex
                gc.collect()
                if use_gpu:
                    blob_vkallocator.clear()
                    staging_vkallocator.clear()
                # TODO: Have someone running into this issue enable this and see if it fixes anything
                # ncnn.destroy_gpu_instance()
                raise RuntimeError(
                    "A critical error has occurred. You may need to restart chaiNNer in order for NCNN upscaling to start working again."
                ) from e
            # Check to see if its actually the NCNN out of memory error
            if "failed" in str(e):
                # clear VRAM
                logger.debug("NCNN out of VRAM, clearing VRAM and splitting.")
                ex = None
                del ex
                gc.collect()
                if use_gpu:
                    blob_vkallocator.clear()
                    staging_vkallocator.clear()
                return Split()
            else:
                # Re-raise the exception if not an OOM error
                raise

    return auto_split(img, upscale, tiler)

nodes/impl/ncnn/auto_split.py
vladmandic/sd-extension-chainner
import os
from copy import deepcopy
from io import BufferedReader, StringIO
from json import load as jload
from typing import Dict, List, Tuple, Union

import numpy as np
from nodes.log import logger

from ...utils.checked_cast import checked_cast

param_schema_file = os.path.join(
    os.path.dirname(os.path.realpath(__file__)), "param_schema.json"
)
with open(param_schema_file, encoding="utf-8") as schemaf:
    param_schema = jload(schemaf)

DTYPE_FP32 = b"\x00\x00\x00\x00"
DTYPE_FP16 = b"\x47\x6b\x30\x01"
DTYPE_DICT = {b"\x00\x00\x00\x00": np.float32, b"\x47\x6b\x30\x01": np.float16}


class UnaryOpTypes:
    ABS = 0
    NEG = 1
    FLOOR = 2
    CEIL = 3
    SQUARE = 4
    SQRT = 5
    RSQ = 6
    EXP = 7
    LOG = 8
    SIN = 9
    COS = 10
    TAN = 11
    ASIN = 12
    ACOS = 13
    ATAN = 14
    RECIPROCAL = 15
    TANH = 16


class BinaryOpTypes:
    ADD = 0
    SUB = 1
    MUL = 2
    DIV = 3
    MAX = 4
    MIN = 5
    POW = 6
    RSUB = 7
    RDIV = 8


class CastElementTypes:
    AUTO = 0
    FLOAT32 = 1
    FLOAT16 = 2
    INT8 = 3
    BFLOAT16 = 4


class EltwiseOpTypes:
    PROD = 0
    SUM = 1
    MAX = 2


class GruDirectionFlags:
    FORWARD = 0
    REVERSE = 1
    BIDIRECTIONAL = 2


class InterpResizeTypes:
    NEAREST = 1
    BILINEAR = 2
    BICUBIC = 3


class NormalizeEpsModes:
    CAFFE = 0
    PYTORCH = 1
    TENSORFLOW = 2


class PaddingTypes:
    CONSTANT = 0
    REPLICATE = 1
    REFLECT = 2


class PadModes:
    FULL = 0
    VALID = 1
    SAMEUPPER = 2
    SAMELOWER = 3


class PermuteOrderTypes:
    WH_WHC_WHDC = 0
    HW_HWC_HWDC = 1
    WCH_WDHC = 2
    CWH_DWHC = 3
    HCW_HDWC = 4
    CHW_DHWC = 5
    WHCD = 6
    HWCD = 7
    WCHD = 8
    CWHD = 9
    HCWD = 10
    CHWD = 11
    WDCH = 12
    DWCH = 13
    WCDH = 14
    CWDH = 15
    DCWH = 16
    CDWH = 17
    HDCW = 18
    DHCW = 19
    HCDW = 20
    CHDW = 21
    DCHW = 22
    CDHW = 23


class ReductionOpTypes:
    SUM = 0
    ASUM = 1
    SUMSQ = 2
    MEAN = 3
    MAX = 4
    MIN = 5
    PROD = 6
    L1 = 7
    L2 = 8
    LOGSUM = 9
    LOGSUMEXP = 10


class GridSampleSampleTypes:
    NEAREST = 1
    BILINEAR = 2
    BICUBIC = 3


class GridSamplePadModes:
    ZEROS = 1
    BORDER = 2
    REFLECTION = 3


class LrnRegionTypes:
    ACROSS_CHANNELS = 0
    WITH_CHANNEL = 1


class NcnnWeight:
    def __init__(self, weight: np.ndarray, quantize_tag: bytes = b""):
        self.quantize_tag = quantize_tag
        self.weight = weight

    @property
    def shape(self) -> tuple:
        return self.weight.shape


class NcnnParam:
    def __init__(
        self,
        pid: str,
        name: str,
        value: Union[float, int, List[Union[float, int]]],
        default: Union[float, int],
    ) -> None:
        self.id: str = pid
        self.name: str = name
        self.value: Union[float, int, List[Union[float, int]]] = value
        self.default: Union[float, int] = default


class NcnnParamCollection:
    def __init__(
        self,
        op: str,
        param_dict: Union[Dict[int, NcnnParam], None] = None,
    ) -> None:
        self.op: str = op
        self.param_dict: Dict[int, NcnnParam] = {} if param_dict is None else param_dict
        self.weight_order: Dict[str, List[int]] = (
            param_schema[self.op]["weightOrder"] if self.op else {}
        )

    def __getitem__(self, pid: int) -> NcnnParam:
        try:
            return self.param_dict[pid]
        except KeyError as exc:
            idstr = str(pid)
            param_dict = param_schema[self.op]
            try:
                param = param_dict[idstr]
            except KeyError:
                logger.error(f"chaiNNer: op {self.op} does not have param {pid}, please report")
                raise

            defaultValue = param["defaultValue"]
            value = param["defaultValue"]
            if isinstance(value, str):
                for key, val in list(param_dict.items())[:-1]:
                    if value == val["paramPhase"]:
                        try:
                            value = self.param_dict[int(key)].value
                        except KeyError:
                            value = val["defaultValue"]
                        defaultValue = val["defaultValue"]

                        break
                else:
                    msg = f"Op {self.op} does not have param {value}, please report"
                    raise KeyError(msg) from exc

            return NcnnParam(idstr, param["paramPhase"], value, defaultValue)

    def __setitem__(
        self, pid: int, value: Union[float, int, List[Union[float, int]]]
    ) -> None:
        idstr = str(pid)
        param_dict = param_schema[self.op]
        try:
            param = param_dict[idstr]
        except KeyError:
            logger.error(f"chaiNNer: op {self.op} does not have param {idstr}, please report")
            raise
        name = param["paramPhase"]
        def_val = param["defaultValue"]

        self.param_dict[pid] = NcnnParam(idstr, name, value, def_val)

    def __delitem__(self, key: int) -> None:
        try:
            del self.param_dict[key]
        except KeyError:
            pass

    def __contains__(self, item) -> bool:
        if item in self.param_dict:
            return True
        return False

    def __str__(self) -> str:
        output = ""
        param_dict = param_schema[self.op]
        self.param_dict = dict(sorted(self.param_dict.items()))
        for v in self.param_dict.values():
            if v.value == v.default:
                continue
            if isinstance(v.default, str):
                pid = None
                for key, val in list(param_dict.items())[:-1]:
                    if v.default == val["paramPhase"]:
                        pid = int(key)
                        break
                else:
                    msg = f"Op {self.op} does not have param {v.default}, please report"
                    raise KeyError(msg)

                # If a param that defaults to the value of another param, if it's value
                # equals that of the second param or its default, skip writing it
                if (
                    v.value == self.param_dict[pid].value
                    or v.value == self.param_dict[pid].default
                ):
                    continue

            if isinstance(v.value, list):
                output += " -233" + v.id.zfill(2) + "="
            else:
                output += " " + v.id + "="

            if isinstance(v.value, float):
                v_str = np.format_float_scientific(v.value, 6, False, exp_digits=2)
            elif isinstance(v.value, list):
                v_str = ",".join(
                    [
                        np.format_float_scientific(n, 6, False, exp_digits=2)
                        if isinstance(n, float)
                        else str(n)
                        for n in v.value
                    ]
                )
            else:
                v_str = str(v.value)

            output += v_str

        return output

    def set_op(self, op: str) -> None:
        self.op = op
        self.weight_order = param_schema[op]["weightOrder"]


class NcnnLayer:
    def __init__(
        self,
        op_type: str = "",
        name: str = "",
        num_inputs: int = 0,
        num_outputs: int = 0,
        inputs: Union[List[str], None] = None,
        outputs: Union[List[str], None] = None,
        params: Union[NcnnParamCollection, None] = None,
        weight_data: Union[Dict[str, NcnnWeight], None] = None,
    ):
        self.op_type: str = op_type
        self.name: str = name
        self.num_inputs: int = num_inputs
        self.num_outputs: int = num_outputs
        self.inputs: List[str] = [] if inputs is None else inputs
        self.outputs: List[str] = [] if outputs is None else outputs
        self.params: NcnnParamCollection = (
            NcnnParamCollection(op_type) if params is None else params
        )
        self.weight_data: Dict[str, NcnnWeight] = (
            {} if weight_data is None else weight_data
        )

    def add_param(
        self, pid: int, value: Union[float, int, List[Union[float, int]]]
    ) -> None:
        self.params[pid] = value

    def add_weight(
        self,
        weight_name: str,
        data: Union[float, int, np.ndarray],
        quantize_tag: bytes = b"",
    ) -> int:
        if isinstance(data, float):
            data_array = np.array(data, np.float32)
        elif isinstance(data, int):
            data_array = np.array(data, np.int32)
        else:
            data_array = data

        if quantize_tag == DTYPE_FP16:
            data_array = data_array.astype(np.float16)
        else:
            # Since int8 not supported, all data that is not fp16 is fp32.
            # This covers issues caused by converting fp16 ONNX models.
            data_array = data_array.astype(np.float32)
        self.weight_data[weight_name] = NcnnWeight(data_array, quantize_tag)

        return len(quantize_tag) + len(data_array.tobytes())


class NcnnModel:
    def __init__(
        self,
        node_count: int = 0,
        blob_count: int = 0,
    ) -> None:
        self.node_count: int = node_count
        self.blob_count: int = blob_count
        self.layers: List[NcnnLayer] = []
        self.bin_length = 0

    @property
    def magic(self):
        return "7767517"

    @staticmethod
    def load_from_file(param_path: str = "", bin_path: str = "") -> "NcnnModel":
        if bin_path == "":
            bin_path = param_path.replace(".param", ".bin")
        elif param_path == "":
            param_path = bin_path.replace(".bin", ".param")

        model = NcnnModel()
        with open(param_path, "r", encoding="utf-8") as paramf:
            with open(bin_path, "rb") as binf:
                paramf.readline()
                counts = paramf.readline().strip().split(" ")
                model.node_count = int(counts[0])
                model.blob_count = int(counts[1])

                for line in paramf:
                    op_type, layer = model.parse_param_layer(line)
                    layer.weight_data = model.load_layer_weights(binf, op_type, layer)
                    model.add_layer(layer)

                binf.seek(0, os.SEEK_END)
                model.bin_length = binf.tell()

        return model

    @staticmethod
    def interp_layers(
        a: NcnnLayer, b: NcnnLayer, alpha_a: float
    ) -> Tuple[NcnnLayer, bytes]:
        weights_a = a.weight_data
        weights_b = b.weight_data
        weights_interp: Dict[str, NcnnWeight] = {}
        layer_bytes = b""

        if weights_a:
            assert len(weights_a) == len(
                weights_b
            ), "All corresponding nodes must have same number of weights"

            layer_bytes_list = []
            for weight_name, weight_a in weights_a.items():
                try:
                    weight_b = weights_b[weight_name]
                except KeyError:
                    logger.error(f"chaiNNer: weights in node {a.name} and {b.name} do not correspond")
                    raise

                assert (
                    weight_a.shape == weight_b.shape
                ), "Corresponding weights must have the same size and shape"

                assert len(weight_a.quantize_tag) == len(
                    weight_b.quantize_tag
                ), "Weights must either both have or both not have a quantize tag"

                if (
                    weight_a.quantize_tag == DTYPE_FP16
                    and weight_b.quantize_tag == DTYPE_FP32
                ):
                    weight_b.quantize_tag = DTYPE_FP16
                    weight_b.weight = weight_b.weight.astype(np.float16)
                elif (
                    weight_a.quantize_tag == DTYPE_FP32
                    and weight_b.quantize_tag == DTYPE_FP16
                ):
                    weight_a.quantize_tag = DTYPE_FP16
                    weight_a.weight = weight_a.weight.astype(np.float16)

                weight_c = NcnnWeight(
                    (weight_a.weight * alpha_a + weight_b.weight * (1 - alpha_a)),
                    weight_a.quantize_tag,
                )
                layer_bytes_list.append(
                    weight_c.quantize_tag + weight_c.weight.tobytes()
                )

                weights_interp[weight_name] = weight_c

            layer_bytes = b"".join(layer_bytes_list)

        return (
            NcnnLayer(
                a.op_type,
                a.name,
                a.num_inputs,
                a.num_outputs,
                a.inputs,
                a.outputs,
                a.params,
                weights_interp,
            ),
            layer_bytes,
        )

    def add_layer(self, layer: NcnnLayer) -> None:
        self.layers.append(layer)

    def parse_param_layer(self, layer_str: str) -> Tuple[str, NcnnLayer]:
        param_list = layer_str.strip().split()
        op_type, name = param_list[:2]
        assert op_type != "MemoryData", "This NCNN param file contains invalid layers"

        num_inputs = int(param_list[2])
        num_outputs = int(param_list[3])
        input_end = 4 + num_inputs
        output_end = input_end + num_outputs
        inputs = list(param_list[4:input_end])
        outputs = list(param_list[input_end:output_end])

        params = param_list[output_end:]
        param_dict = {}
        for param_str in params:
            ks, vs = param_str.split("=")
            k = int(ks)
            if k < 0:
                v = []
                for vi in vs.split(","):
                    vi = float(vi) if "." in vi or "e" in vi else int(vi)
                    v.append(vi)
                k = abs(k + 23300)
                ks = str(k)
            elif "." in vs or "e" in vs:
                v = float(vs)
            else:
                v = int(vs)

            param = NcnnParam(
                ks,
                param_schema[op_type][ks]["paramPhase"],
                v,
                param_schema[op_type][ks]["defaultValue"],
            )
            param_dict[k] = param

        return op_type, NcnnLayer(
            op_type,
            name,
            num_inputs,
            num_outputs,
            inputs,
            outputs,
            NcnnParamCollection(op_type, param_dict),
        )

    def load_layer_weights(
        self, binf: BufferedReader, op_type: str, layer: NcnnLayer
    ) -> Dict[str, NcnnWeight]:
        weight_dict = {}
        if op_type == "BatchNorm":
            channels_data = checked_cast(int, layer.params[0].value) * 4
            slope = np.frombuffer(binf.read(channels_data), np.float32)
            weight_dict["slope"] = NcnnWeight(slope)
            mean = np.frombuffer(binf.read(channels_data), np.float32)
            weight_dict["mean"] = NcnnWeight(mean)
            variance = np.frombuffer(binf.read(channels_data), np.float32)
            weight_dict["variance"] = NcnnWeight(variance)
            bias = np.frombuffer(binf.read(channels_data), np.float32)
            weight_dict["bias"] = NcnnWeight(bias)
        elif op_type in ("Convolution", "ConvolutionDepthWise"):
            quantize_tag = binf.read(4)
            dtype = DTYPE_DICT[quantize_tag]
            weight_data_length = checked_cast(int, layer.params[6].value)
            weight_data_size = (
                weight_data_length * 2
                if quantize_tag == DTYPE_FP16
                else weight_data_length * 4
            )

            has_bias = layer.params[5].value

            num_filters = checked_cast(int, layer.params[0].value)
            kernel_w = checked_cast(int, layer.params[1].value)
            kernel_h = checked_cast(int, layer.params[11].value)
            if op_type == "ConvolutionDepthWise":
                group = checked_cast(int, layer.params[7].value)
                num_input = (
                    weight_data_length // (num_filters // group) // kernel_w // kernel_h
                )
                shape = (
                    group,
                    num_filters // group,
                    num_input // group,
                    kernel_h,
                    kernel_w,
                )
            else:
                num_input = weight_data_length // num_filters // kernel_w // kernel_h
                shape = (num_filters, num_input, kernel_h, kernel_w)

            weight_data = np.frombuffer(binf.read(weight_data_size), dtype)
            weight_data = weight_data.reshape(shape)
            weight_dict["weight"] = NcnnWeight(weight_data, quantize_tag)

            if has_bias:
                bias_data_size = num_filters * 4
                bias_data = np.frombuffer(binf.read(bias_data_size), np.float32)
                weight_dict["bias"] = NcnnWeight(bias_data)
        elif op_type == "Deconvolution":
            quantize_tag = binf.read(4)
            dtype = DTYPE_DICT[quantize_tag]
            weight_data_length = checked_cast(int, layer.params[6].value)
            weight_data_size = (
                weight_data_length * 2
                if quantize_tag == DTYPE_FP16
                else weight_data_length * 4
            )

            has_bias = layer.params[5].value

            num_filters = checked_cast(int, layer.params[0].value)
            kernel_w = checked_cast(int, layer.params[1].value)
            kernel_h = checked_cast(int, layer.params[11].value)
            num_input = weight_data_length // num_filters // kernel_w // kernel_h
            shape = (num_filters, num_input, kernel_h, kernel_w)

            weight_data = np.frombuffer(binf.read(weight_data_size), dtype)
            weight_data = weight_data.reshape(shape)
            weight_dict["weight"] = NcnnWeight(weight_data, quantize_tag)

            if has_bias:
                bias_data_size = num_filters * 4
                bias_data = np.frombuffer(binf.read(bias_data_size), np.float32)
                weight_dict["bias"] = NcnnWeight(bias_data)
        elif op_type == "InnerProduct":
            quantize_tag = binf.read(4)
            dtype = DTYPE_DICT[quantize_tag]
            weight_data_length = layer.params[2].value
            assert isinstance(weight_data_length, int), "Weight data size must be int"
            weight_data_size = (
                weight_data_length * 2
                if quantize_tag == DTYPE_FP16
                else weight_data_length * 4
            )
            weight_data = np.frombuffer(binf.read(weight_data_size), dtype)
            num_output = layer.params[0].value
            assert isinstance(num_output, int), "Num output must be int"
            num_input = weight_data_length // num_output
            weight_data = weight_data.reshape((num_input, num_output))
            weight_dict["weight"] = NcnnWeight(weight_data, quantize_tag)

            has_bias = layer.params[1].value
            if has_bias == 1:
                bias_data_size = num_output * 4
                bias_data = np.frombuffer(binf.read(bias_data_size), np.float32)
                weight_dict["bias"] = NcnnWeight(bias_data)
        elif op_type == "PReLU":
            num_slope = layer.params[0].value
            assert isinstance(num_slope, int), "Num slopes must be int"
            slope_data_size = num_slope * 4
            slope_data = np.frombuffer(binf.read(slope_data_size), np.float32)
            weight_dict["slope"] = NcnnWeight(slope_data)
        elif op_type == "Scale":
            scale_data_length = layer.params[0].value
            assert isinstance(scale_data_length, int), "Scale data size must be int"
            if scale_data_length != -233:
                quantize_tag = binf.read(4)
                dtype = DTYPE_DICT[quantize_tag]
                scale_data_size = (
                    scale_data_length * 2
                    if quantize_tag == DTYPE_FP16
                    else scale_data_length * 4
                )
                scale_data = np.frombuffer(binf.read(scale_data_size), dtype)
                weight_dict["weight"] = NcnnWeight(scale_data, quantize_tag)

                has_bias = layer.params[1].value
                if has_bias == 1:
                    bias_data = np.frombuffer(
                        binf.read(scale_data_length * 4), np.float32
                    )
                    weight_dict["bias"] = NcnnWeight(bias_data)
        else:
            if len(layer.params.weight_order) != 0:
                error_msg = f"Load weights not added for {op_type} yet, please report"
                raise ValueError(error_msg)

        return weight_dict

    def write_param(self, filename: str = "") -> str:
        with StringIO() as p:
            p.write(f"{self.magic}\n{self.node_count} {self.blob_count}\n")

            for layer in self.layers:
                if layer.op_type == "ncnnfused":
                    continue

                p.write(
                    f"{layer.op_type:<16}"
                    f" {layer.name:<24}"
                    f" {layer.num_inputs}"
                    f" {layer.num_outputs}"
                )
                if layer.inputs:
                    p.write(f" {' '.join(layer.inputs)}")
                if layer.outputs:
                    p.write(f" {' '.join(layer.outputs)}")
                if layer.params.param_dict:
                    param_str = str(layer.params)
                    if param_str:
                        p.write(f"{param_str}")
                p.write("\n")

            if filename:
                with open(filename, "w", encoding="utf-8") as f:
                    f.write(p.getvalue())
                return ""
            else:
                return p.getvalue()

    def serialize_weights(self) -> bytes:
        layer_weights = [
            b"".join((w.quantize_tag, np.ndarray.tobytes(w.weight)))
            for l in self.layers
            for w in l.weight_data.values()
            if l.weight_data and l.op_type != "ncnnfused"
        ]

        return b"".join(layer_weights)

    def write_bin(self, filename: str) -> None:
        with open(filename, "wb") as f:
            f.write(self.serialize_weights())

    def interpolate(self, model_b: "NcnnModel", alpha: float) -> "NcnnModel":
        interp_model = deepcopy(self)

        layer_a_weights = [(i, l) for i, l in enumerate(self.layers) if l.weight_data]
        layer_b_weights = [
            (i, l) for i, l in enumerate(model_b.layers) if l.weight_data
        ]

        assert len(layer_a_weights) == len(
            layer_b_weights
        ), "Models must have same number of layers containing weights"

        weight_bytes_list = []
        for layer_a, layer_b in zip(layer_a_weights, layer_b_weights):
            interp_layer, layer_bytes = NcnnModel.interp_layers(
                layer_a[1], layer_b[1], alpha
            )
            interp_model.layers[layer_a[0]] = interp_layer
            weight_bytes_list.append(layer_bytes)

        return interp_model

    @property
    def bin(self) -> bytes:
        return self.serialize_weights()


class NcnnModelWrapper:
    def __init__(self, model: NcnnModel) -> None:
        self.model: NcnnModel = model
        scale, in_nc, out_nc, nf, fp = NcnnModelWrapper.get_broadcast_data(model)
        self.scale: int = scale
        self.nf: int = nf
        self.in_nc: int = in_nc
        self.out_nc: int = out_nc
        self.fp: str = fp

    @staticmethod
    def get_broadcast_data(model: NcnnModel) -> Tuple[int, int, int, int, str]:
        scale = 1.0
        in_nc = 0
        out_nc = 0
        nf = 0
        fp = "fp32"
        pixel_shuffle = 1
        found_first_conv = False
        current_conv = None

        for i, layer in enumerate(model.layers):
            if layer.op_type == "Interp":
                try:
                    if (
                        model.layers[i + 1].op_type != "BinaryOp"
                        and model.layers[i + 1].params[0].value != 0
                    ):
                        scale *= checked_cast(float, layer.params[1].value)
                except IndexError:
                    scale *= checked_cast(float, layer.params[1].value)
            elif layer.op_type == "PixelShuffle":
                scale *= checked_cast(int, layer.params[0].value)
                pixel_shuffle *= checked_cast(int, layer.params[0].value)
            elif layer.op_type in (
                "Convolution",
                "Convolution1D",
                "ConvolutionDepthWise",
            ):
                if found_first_conv is not True:
                    nf, in_nc = NcnnModelWrapper.get_nf_and_in_nc(layer)
                    if layer.weight_data["weight"].quantize_tag == DTYPE_FP16:
                        fp = "fp16"
                    found_first_conv = True

                scale /= checked_cast(int, layer.params[3].value)
                current_conv = layer
            elif layer.op_type in ("Deconvolution", "DeconvolutionDepthWise"):
                if found_first_conv is not True:
                    nf, in_nc = NcnnModelWrapper.get_nf_and_in_nc(layer)
                    found_first_conv = True

                scale *= checked_cast(int, layer.params[3].value)
                current_conv = layer

        assert (
            current_conv is not None
        ), "Cannot broadcast; model has no Convolution layers"

        out_nc = checked_cast(int, current_conv.params[0].value) // pixel_shuffle**2

        assert scale >= 1, "Models with scale less than 1x not supported"
        assert scale % 1 == 0, f"Model not supported, scale {scale} is not an integer"

        return int(scale), in_nc, out_nc, nf, fp

    @staticmethod
    def get_nf_and_in_nc(layer: NcnnLayer) -> Tuple[int, int]:
        nf = layer.params[0].value
        kernel_w = layer.params[1].value
        try:
            kernel_h = layer.params[11].value
        except KeyError:
            kernel_h = kernel_w
        weight_data_size = layer.params[6].value

        assert (
            isinstance(nf, int)
            and isinstance(kernel_w, int)
            and isinstance(kernel_h, int)
            and isinstance(weight_data_size, int)
        ), "Out nc, kernel width and height, and weight data size must all be ints"
        in_nc = weight_data_size // nf // kernel_w // kernel_h

        return nf, in_nc

nodes/impl/ncnn/model.py
vladmandic/sd-extension-chainner
import numpy as np

from ...utils.checked_cast import checked_cast
from .model import BinaryOpTypes as BOT
from .model import EltwiseOpTypes as EOT
from .model import NcnnLayer, NcnnModel


class NcnnOptimizer:
    def __init__(self, model: NcnnModel) -> None:
        self.model = model

    def __fuse_batchnorm_scale(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "BatchNorm":
                # BatchNorm - Scale
                batchnorm_output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "Scale":
                        continue
                    if len(self.model.layers[j].inputs) != 1:
                        continue
                    if self.model.layers[j].inputs[0] == batchnorm_output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                # fuse BatchNorm - Scale to BatchNorm
                scale = self.model.layers[j]

                bias = layer.weight_data["bias"].weight

                layer.weight_data["slope"].weight = (
                    layer.weight_data["slope"].weight
                    * scale.weight_data["scale"].weight
                )

                bias = bias * scale.weight_data["scale"].weight
                if scale.params[1].value:
                    bias += scale.weight_data["bias"].weight
                layer.weight_data["bias"].weight = bias

                layer.outputs[0] = scale.outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                scale.op_type = "ncnnfused"

    def __fuse_x_batchnorm(self):
        """Combines fuse_convolution_batchnorm, fuse_convolutiondepthwise_batchnorm,
        fuse_deconvolution_batchnorm, fuse_deconvolutiondepthwise_batchnorm, and
        fuse_innerproduct_batchnorm"""

        for i, layer in enumerate(self.model.layers):
            if layer.op_type in (
                "Convolution",
                "ConvolutionDepthWise",
                "Deconvolution",
                "DeconvolutionDepthWise",
                "InnerProduct",
            ):
                # Convolution - BatchNorm
                conv_output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "BatchNorm":
                        continue
                    if len(self.model.layers[j].inputs) != 1:
                        continue
                    if self.model.layers[j].inputs[0] == conv_output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                # fuse Convolution - BatchNorm to Convolution
                batchnorm = self.model.layers[j]

                channels = checked_cast(int, batchnorm.params[0].value)
                eps = checked_cast(float, batchnorm.params[1].value)

                # a = bias - slope * mean / sqrt(var + eps)
                # b = slope / sqrt(var + eps)
                # value = value * b + a
                a = np.ndarray((channels,))
                b = np.ndarray((channels,))
                sqrt_var = np.sqrt(batchnorm.weight_data["variance"].weight + eps)
                a = (
                    batchnorm.weight_data["bias"].weight
                    - batchnorm.weight_data["slope"].weight
                    * batchnorm.weight_data["mean"].weight
                    / sqrt_var
                )
                b = batchnorm.weight_data["slope"].weight / sqrt_var

                bias_term = 1 if layer.op_type == "InnerProduct" else 5

                if layer.params[bias_term].value == 0:
                    # init bias as zero
                    layer.params[bias_term] = 1
                    layer.add_weight("bias", np.zeros(channels, dtype=np.float32))

                weight = layer.weight_data["weight"].weight
                layer.weight_data["weight"].weight = weight * (
                    np.transpose(
                        np.broadcast_to(b, weight.shape[::-1]).astype(weight.dtype),
                        (3, 2, 1, 0),
                    )
                )

                layer.weight_data["bias"].weight = (
                    layer.weight_data["bias"].weight * b + a
                )

                self.model.layers[i].outputs[0] = self.model.layers[j].outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                batchnorm.op_type = "ncnnfused"

    def __fuse_x_mul(self):
        """Combines fuse_convolution_mul, fuse_convolutiondepthwise_mul,
        and fuse_deconvolution_mul"""

        for i, layer in enumerate(self.model.layers):
            if layer.op_type in (
                "Convolution",
                "ConvolutionDepthWise",
                "Deconvolution",
            ):
                # Convolution - BinaryOp
                output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "BinaryOp":
                        continue
                    if self.model.layers[j].num_inputs != 2:
                        continue
                    if self.model.layers[j].inputs[0] == output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                # fuse Convolution - BinaryOp to Convolution
                binaryop = self.model.layers[j]

                if binaryop.params[0].value != BOT.MUL or binaryop.params[1].value:
                    continue

                # MemoryData - ..... - BinaryOp
                k = 0
                for k in range(j):
                    if self.model.layers[k].op_type != "MemoryData":
                        continue
                    if self.model.layers[k].outputs[0] == binaryop.inputs[1]:
                        break
                else:
                    k += 1

                if k == j:
                    continue

                memorydata = self.model.layers[k]

                channels = checked_cast(int, layer.params[0].value)

                if (
                    memorydata.params[0].value != channels
                    or memorydata.params[1].value != 0
                    or memorydata.params[2].value != 0
                ):
                    # not bias-like broadcasting type
                    continue

                data = memorydata.weight_data["data"].weight

                weight = layer.weight_data["weight"].weight
                layer.weight_data["weight"].weight = weight * (
                    np.transpose(
                        np.broadcast_to(data, weight.shape[::-1]).astype(weight.dtype),
                        (3, 2, 1, 0),
                    )
                )

                try:
                    layer.weight_data["bias"].weight = (
                        layer.weight_data["bias"].weight * data
                    )
                except KeyError:
                    pass

                self.model.layers[i].outputs[0] = self.model.layers[j].outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                binaryop.op_type = "ncnnfused"

    def __fuse_x_add(self):
        """Combines fuse_convolution_add, fuse_convolutiondepthwise_add,
        fuse_deconvolution_add, and fuse_innerproduct_add"""

        for i, layer in enumerate(self.model.layers):
            if layer.op_type in (
                "Convolution",
                "ConvolutionDepthWise",
                "Deconvolution",
                "InnerProduct",
            ):
                # Convolution - Add
                output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "BinaryOp":
                        continue
                    if self.model.layers[j].num_inputs != 2:
                        continue
                    if self.model.layers[j].inputs[0] == output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                # fuse Convolution - BinaryOp to Convolution
                binaryop = self.model.layers[j]

                if binaryop.params[0].value != BOT.ADD or binaryop.params[1].value:
                    continue

                # MemoryData - ..... - BinaryOp
                k = 0
                for k in range(j):
                    if self.model.layers[k].op_type != "MemoryData":
                        continue
                    if self.model.layers[k].outputs[0] == binaryop.inputs[1]:
                        break
                else:
                    k += 1

                if k == j:
                    continue

                memorydata = self.model.layers[k]

                channels = checked_cast(int, layer.params[0].value)

                if not (
                    memorydata.params[0].value == channels
                    and memorydata.params[1].value == 0
                    and memorydata.params[2].value == 0
                ) or (
                    memorydata.params[0].value == 1
                    and memorydata.params[1].value == 1
                    and memorydata.params[2].value == channels
                ):
                    # not bias-like broadcasting type
                    continue

                bias_term = 1 if layer.op_type == "InnerProduct" else 5
                bias_data = memorydata.weight_data["data"].weight.reshape(channels)

                if layer.params[bias_term].value == 0:
                    # init bias
                    layer.params[bias_term] = 1
                    layer.add_weight("bias", bias_data)
                else:
                    layer.weight_data["bias"].weight = (
                        layer.weight_data["bias"].weight + bias_data
                    )

                self.model.layers[i].outputs[0] = self.model.layers[j].outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                binaryop.op_type = "ncnnfused"

    def __fuse_innerproduct_dropout(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "InnerProduct":
                # InnerProduct - Dropout
                output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "Dropout":
                        continue
                    if self.model.layers[j].num_inputs != 1:
                        continue
                    if self.model.layers[j].inputs[0] == output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                # fuse InnerProduct - Dropout to InnerProduct
                dropout = self.model.layers[j]

                scale = checked_cast(float, dropout.params[0].value)
                if scale != 1:
                    layer.weight_data["weight"].weight = (
                        layer.weight_data["weight"].weight * scale
                    )

                    if layer.params[1].value == 1:
                        layer.weight_data["bias"].weight = (
                            layer.weight_data["bias"].weight * scale
                        )

                self.model.layers[i].outputs[0] = self.model.layers[j].outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                dropout.op_type = "ncnnfused"

    def __fuse_x_activation(self):
        """Combines fuse_convolution_activation, fuse_convolution1d_activation,
        fuse_convolutiondepthwise_activation, fuse_deconvolution_activation,
        fuse_deconvolutiondepthwise_activation, and fuse_innerproduct_activation"""

        for i, layer in enumerate(self.model.layers):
            if layer.op_type in (
                "Convolution",
                "Convolution1D",
                "ConvolutionDepthWise",
                "Deconvolution",
                "DeconvolutionDepthWise",
                "InnerProduct",
            ):
                # Convolution - Activation
                output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type not in (
                        "ReLU",
                        "Clip",
                        "Sigmoid",
                        "Mish",
                        "Hardswish",
                    ):
                        continue
                    if (
                        self.model.layers[j].op_type == "Mish"
                        and layer.op_type in ("Deconvolution", "DeconvolutionDepthWise")
                    ) or (
                        self.model.layers[j].op_type == "HardSwish"
                        and layer.op_type
                        in (
                            "Convolution1D",
                            "Deconvolution",
                            "DeconvolutionDepthWise",
                        )
                    ):
                        continue
                    if self.model.layers[j].num_inputs != 1:
                        continue
                    if self.model.layers[j].inputs[0] == output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                # fuse Convolution - Activation to Convolution
                act = self.model.layers[j]

                if act.op_type == "ReLU":
                    if act.params[0].value == 0:
                        layer.params[9] = 1
                    else:
                        layer.params[9] = 2
                        layer.params[10] = [1, checked_cast(float, act.params[0].value)]
                elif act.op_type == "Clip":
                    layer.params[9] = 3
                    layer.params[10] = [
                        2,
                        checked_cast(float, act.params[0].value),
                        checked_cast(float, act.params[1].value),
                    ]
                elif act.op_type == "Sigmoid":
                    layer.params[9] = 4
                elif act.op_type == "Mish":
                    layer.params[9] = 5
                elif act.op_type == "HardSwish":
                    layer.params[9] = 6
                    layer.params[10] = [
                        2,
                        checked_cast(float, act.params[0].value),
                        checked_cast(float, act.params[1].value),
                    ]

                self.model.layers[i].outputs[0] = self.model.layers[j].outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                act.op_type = "ncnnfused"

    def __fuse_memorydata_binaryop(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "MemoryData":
                # MemoryData - BinaryOp
                output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "BinaryOp":
                        continue
                    if self.model.layers[j].num_inputs != 2:
                        continue
                    if (
                        self.model.layers[j].inputs[0] == output
                        or self.model.layers[j].inputs[1] == output
                    ):
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                # fuse MemoryData - BinaryOp to BinaryOp
                binaryop = self.model.layers[j]

                if (
                    layer.params[0].value != 1
                    or layer.params[1].value != 0
                    or layer.params[2].value != 0
                ):
                    # not a scalar
                    continue

                memorydata_index = 1
                if binaryop.inputs[0] == output:
                    op_type = checked_cast(int, binaryop.params[0].value)
                    if op_type == BOT.ADD:
                        memorydata_index = 0
                    elif op_type == BOT.SUB:
                        binaryop.params[0] = BOT.RSUB
                        memorydata_index = 0
                    elif op_type == BOT.DIV:
                        binaryop.params[0] = BOT.RDIV
                        memorydata_index = 0
                    else:
                        # non-interchangeable binaryop
                        continue

                binaryop.params[1] = 1
                binaryop.params[2] = layer.weight_data["data"].weight[0]

                binaryop.inputs.pop(memorydata_index)
                self.model.node_count -= 1
                self.model.blob_count -= 1
                layer.op_type = "ncnnfused"

        i = 0
        while i in range(len(self.model.layers)):
            if self.model.layers[i].op_type != "MemoryData":
                # MemoryData - Split - BinaryOp
                output = self.model.layers[i].outputs[0]

                j0 = i
                for j0 in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j0].op_type != "Split":
                        continue
                    if self.model.layers[j0].num_inputs != 1:
                        continue
                    if self.model.layers[j0].inputs[0] == output:
                        break
                else:
                    j0 += 1

                if j0 == len(self.model.layers):
                    i += 1
                    continue

                split_output_index = -1
                j1 = i
                for j1 in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j1].op_type != "BinaryOp":
                        continue
                    if self.model.layers[j1].num_inputs != 2:
                        continue
                    for k in range(self.model.layers[j0].num_outputs):
                        if (
                            self.model.layers[j1].inputs[0]
                            == self.model.layers[j0].outputs[k]
                            or self.model.layers[j1].inputs[1]
                            == self.model.layers[j0].outputs[k]
                        ):
                            split_output_index = k
                            break
                    if split_output_index != -1:
                        break
                else:
                    j1 += 1

                if j1 == len(self.model.layers):
                    i += 1
                    continue

                # fuse MemoryData - Split - BinaryOp to BinaryOp
                split = self.model.layers[j0]
                binaryop = self.model.layers[j1]

                if (
                    self.model.layers[i].params[0].value != 1
                    or self.model.layers[i].params[1].value != 0
                    or self.model.layers[i].params[2].value != 0
                ):
                    # not a scalar
                    i += 1
                    continue

                memorydata_index = 1
                if binaryop.inputs[0] == split.outputs[split_output_index]:
                    op_type = checked_cast(int, binaryop.params[0].value)
                    if op_type in (BOT.ADD, BOT.MUL, BOT.MAX, BOT.MIN):
                        memorydata_index = 0
                    elif op_type == BOT.SUB:
                        binaryop.params[0] = BOT.RSUB
                        memorydata_index = 0
                    elif op_type == BOT.DIV:
                        binaryop.params[0] = BOT.RDIV
                        memorydata_index = 0
                    else:
                        # non-interchangeable binaryop
                        i += 1
                        continue

                binaryop.params[1] = 1
                binaryop.params[2] = self.model.layers[i].weight_data["data"].weight[0]

                binaryop.inputs.pop(memorydata_index)
                binaryop.num_inputs -= 1
                split.outputs.pop(split_output_index)
                split.num_outputs -= 1
                if split.num_outputs == 0:
                    self.model.node_count -= 2
                    self.model.blob_count -= 2
                    split.op_type = "ncnnfused"
                    self.model.layers[i].op_type = "ncnnfused"

                i -= 1

            i += 1

    def __fuse_binaryop_eltwise(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "BinaryOp":
                if layer.num_inputs != 2:
                    continue
                if layer.params[0].value != BOT.ADD or layer.params[1].value:
                    continue

                # BinaryOp - BinaryOp - BinaryOp
                input0 = layer.inputs[0]
                input1 = layer.inputs[1]

                j0 = 0
                for j0 in range(i):
                    if self.model.layers[j0].op_type != "BinaryOp":
                        continue
                    if self.model.layers[j0].num_inputs != 1:
                        continue
                    if self.model.layers[j0].params[0].value != BOT.MUL:
                        continue
                    if self.model.layers[j0].outputs[0] == input0:
                        break
                else:
                    j0 += 1

                j1 = 0
                for j1 in range(i):
                    if self.model.layers[j1].op_type != "BinaryOp":
                        continue
                    if self.model.layers[j1].num_inputs != 1:
                        continue
                    if self.model.layers[j1].params[0].value != BOT.MUL:
                        continue
                    if self.model.layers[j1].outputs[0] == input1:
                        break
                else:
                    j1 += 1

                if j0 == i and j1 == i:
                    continue

                binaryop0 = self.model.layers[j0]
                binaryop1 = self.model.layers[j1]

                eltwise = NcnnLayer(
                    "Eltwise",
                    layer.name,
                    layer.num_inputs,
                    layer.num_outputs,
                    layer.inputs,
                    layer.outputs,
                )
                eltwise.add_param(0, EOT.SUM)
                if j0 != i and j1 != i:
                    # fuse BinaryOp - BinaryOp - BinaryOp to Eltwise
                    eltwise.add_param(
                        1,
                        [
                            2,
                            checked_cast(float, binaryop0.params[2].value),
                            checked_cast(float, binaryop1.params[2].value),
                        ],
                    )
                    eltwise.inputs[0] = binaryop0.inputs[0]
                    eltwise.inputs[1] = binaryop1.inputs[0]
                    self.model.node_count -= 2
                    self.model.blob_count -= 2
                    binaryop0.op_type = "ncnnfused"
                    binaryop1.op_type = "ncnnfused"
                elif j0 != i and j1 == i:
                    # fuse BinaryOp - X - BinaryOp to Eltwise
                    eltwise.add_param(
                        1, [2, checked_cast(float, binaryop0.params[2].value), 1.0]
                    )
                    eltwise.inputs[0] = binaryop0.inputs[0]
                    self.model.node_count -= 1
                    self.model.blob_count -= 1
                    binaryop0.op_type = "ncnnfused"
                else:
                    # fuse X - BinaryOp - BinaryOp to Eltwise
                    eltwise.add_param(
                        1, [2, 1.0, checked_cast(float, binaryop1.params[2].value)]
                    )
                    eltwise.inputs[1] = binaryop1.inputs[0]
                    self.model.node_count -= 1
                    self.model.blob_count -= 1
                    binaryop1.op_type = "ncnnfused"

                self.model.layers[i] = eltwise

    def __eliminate_dropout(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "Dropout":
                if layer.params[0].value != 1:
                    continue

                # Any - Dropout
                dropout_input = layer.inputs[0]

                j = i - 1
                for j in range(i - 1, -1, -1):
                    if self.model.layers[j].op_type == "ncnnfused":
                        continue
                    if self.model.layers[j].num_outputs != 1:
                        continue
                    if self.model.layers[j].outputs[0] == dropout_input:
                        break
                else:
                    j -= 1

                if j == -1:
                    continue

                self.model.layers[j].outputs[0] = layer.outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                layer.op_type = "ncnnfused"

    def __eliminate_pooling1x1(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "Pooling":
                if (
                    layer.params[3].value != 0
                    or layer.params[13].value != 0
                    or layer.params[14].value != 0
                    or layer.params[15].value != 0
                ):
                    continue
                if (
                    layer.params[1].value != 1
                    or layer.params[11].value != 1
                    or layer.params[2].value != 1
                    or layer.params[12].value != 1
                ):
                    continue
                if layer.params[4].value != 0:
                    continue

                # Any - Pooling
                pooling_input = layer.inputs[0]

                top_i = -1
                j = i - 1
                for j in range(i - 1, -1, -1):
                    if self.model.layers[j].op_type == "ncnnfused":
                        continue

                    for k in range(self.model.layers[j].num_outputs):
                        if self.model.layers[j].outputs[k] == pooling_input:
                            top_i = k
                            break

                    if top_i != -1:
                        break
                else:
                    j -= 1

                if j == -1:
                    continue

                self.model.layers[j].outputs[top_i] = layer.outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                layer.op_type = "ncnnfused"

    def __eliminate_noop(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "Noop":
                if layer.num_inputs == 0:
                    # Noop
                    layer.op_type = "ncnnfused"
                    continue

                # Any - Noop
                noop_input = layer.inputs[0]

                j = i - 1
                any_k = -1
                for j in range(i - 1, -1, -1):
                    if self.model.layers[j].op_type == "ncnnfused":
                        continue

                    link_noop = False
                    for k in range(self.model.layers[j].num_outputs):
                        if self.model.layers[j].outputs[k] == noop_input:
                            link_noop = True
                            any_k = k
                            break

                    if link_noop:
                        break
                else:
                    j -= 1

                if j == -1 or any_k == -1:
                    continue

                self.model.layers[j].outputs[any_k] = layer.outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                layer.op_type = "ncnnfused"

    def __eliminate_split(self):
        blob_input_references = []
        for i, layer in enumerate(self.model.layers):
            for input_name in layer.inputs:
                blob_input_references.append(input_name)

        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "Split":
                real_split_output_count = 0
                real_split_output_index = -1
                for j in range(layer.num_outputs):
                    if layer.outputs[j] in blob_input_references:
                        real_split_output_count += 1
                        real_split_output_index = j

                if real_split_output_count > 1:
                    continue

                # Any - Pooling
                split_input = layer.inputs[0]

                top_i = -1
                j = i - 1
                for j in range(i - 1, -1, -1):
                    if self.model.layers[j].op_type == "ncnnfused":
                        continue

                    for k in range(self.model.layers[j].num_outputs):
                        if self.model.layers[j].outputs[k] == split_input:
                            top_i = k
                            break

                    if top_i != -1:
                        break
                else:
                    j -= 1

                if j == -1:
                    continue

                self.model.layers[j].outputs[top_i] = layer.outputs[
                    real_split_output_index
                ]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                layer.op_type = "ncnnfused"

    def __eliminate_orphaned_memorydata(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "MemoryData":
                # MemoryData - X
                memdata_output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type == "ncnnfused":
                        continue

                    orphaned = True
                    for k in range(self.model.layers[j].num_inputs):
                        if self.model.layers[j].inputs[k] == memdata_output:
                            orphaned = False
                            break

                    if not orphaned:
                        break

                if j < len(self.model.layers):
                    continue

                self.model.node_count -= 1
                self.model.blob_count -= 1
                layer.op_type = "ncnnfused"

    def __eliminate_reshape_after_global_pooling(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "Pooling":
                if layer.params[4].value == 0:
                    continue

                # Pooling - Reshape
                pooling_output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "Reshape":
                        continue
                    if self.model.layers[j].num_inputs != 1:
                        continue
                    if self.model.layers[j].inputs[0] == pooling_output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                reshape = self.model.layers[j]

                if (
                    reshape.params[1].value != -233
                    or reshape.params[2].value != -233
                    or reshape.params[3].value != 0
                ):
                    continue

                layer.outputs[0] = reshape.outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                reshape.op_type = "ncnnfused"

    def __eliminate_flatten_after_global_pooling(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "Pooling":
                if layer.params[4].value == 0:
                    continue

                # Pooling - Flatten
                pooling_output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "Flatten":
                        continue
                    if self.model.layers[j].num_inputs != 1:
                        continue
                    if self.model.layers[j].inputs[0] == pooling_output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                flatten = self.model.layers[j]

                layer.outputs[0] = flatten.outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                flatten.op_type = "ncnnfused"

    def __eliminate_flatten_after_innerproduct(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "InnerProduct":
                # InnerProduct - Flatten
                inprod_output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "Flatten":
                        continue
                    if self.model.layers[j].num_inputs != 1:
                        continue
                    if self.model.layers[j].inputs[0] == inprod_output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                flatten = self.model.layers[j]

                layer.outputs[0] = flatten.outputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                flatten.op_type = "ncnnfused"

    def __eliminate_reshape_before_binaryop(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "Reshape":
                if (
                    layer.params[0].value != 1
                    or layer.params[1].value != 1
                    or layer.params[3].value != 1
                ):
                    continue

                # Reshape - BinaryOp
                reshape_output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "BinaryOp":
                        continue
                    if self.model.layers[j].num_inputs != 2:
                        continue
                    if (
                        self.model.layers[j].inputs[0] == reshape_output
                        or self.model.layers[j].inputs[1] == reshape_output
                    ):
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                binaryop = self.model.layers[j]

                input_blob_final = layer.inputs[0]
                if binaryop.inputs[0] == reshape_output:
                    binaryop.inputs[0] = input_blob_final
                if binaryop.inputs[1] == reshape_output:
                    binaryop.inputs[1] = input_blob_final
                self.model.node_count -= 1
                self.model.blob_count -= 1
                layer.op_type = "ncnnfused"

    def __replace_reduction_with_global_pooling(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "Reduction":
                if (
                    layer.params[0].value != 3
                    or layer.params[1].value != 0
                    or layer.params[2].value != 1
                ):
                    continue

                axes = checked_cast(list, layer.params[3].value)
                if len(axes) != 1:
                    continue
                if axes[0] != 2 and axes[0] != 3:
                    continue

                # Reduction(2/3) - Reduction(2)
                reduction1_output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "Reduction":
                        continue
                    if self.model.layers[j].num_inputs != 1:
                        continue
                    if self.model.layers[j].inputs[0] == reduction1_output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                reduction2 = self.model.layers[j]

                if (
                    reduction2.params[0].value != 3
                    or reduction2.params[1].value != 0
                    or reduction2.params[2].value != 1
                ):
                    continue

                axes2 = checked_cast(list, layer.params[3].value)
                if len(axes2) != 1:
                    continue
                if axes2[0] != 2:
                    continue

                pooling = NcnnLayer(
                    "Pooling",
                    reduction2.name,
                    reduction2.num_inputs,
                    reduction2.num_outputs,
                    reduction2.inputs,
                    reduction2.outputs,
                )
                pooling.add_param(0, 1)
                pooling.add_param(4, 1)

                self.model.layers[j] = pooling

                pooling.inputs[0] = layer.inputs[0]
                self.model.node_count -= 1
                self.model.blob_count -= 1
                layer.op_type = "ncnnfused"

    def __replace_prelu_with_leaky_relu(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "PReLU":
                if layer.params[0].value != 1:
                    continue

                relu_layer = NcnnLayer(
                    "ReLU",
                    layer.name,
                    layer.num_inputs,
                    layer.num_outputs,
                    layer.inputs,
                    layer.outputs,
                )
                relu_layer.add_param(
                    0, checked_cast(float, layer.weight_data["slope"].weight[0])
                )

                self.model.layers[i] = relu_layer

    def __replace_convolution_with_innerproduct_after_global_pooling(self):
        for i, layer in enumerate(self.model.layers):
            if layer.op_type == "Pooling":
                if layer.params[4].value == 0:
                    continue

                # Pooling - Convolution
                pooling_output = layer.outputs[0]

                j = i
                for j in range(i + 1, len(self.model.layers)):
                    if self.model.layers[j].op_type != "Convolution":
                        continue
                    if self.model.layers[j].num_inputs != 1:
                        continue
                    if self.model.layers[j].inputs[0] == pooling_output:
                        break
                else:
                    j += 1

                if j == len(self.model.layers):
                    continue

                convolution = self.model.layers[j]

                innerproduct = NcnnLayer(
                    "InnerProduct",
                    convolution.name,
                    convolution.num_inputs,
                    convolution.num_outputs,
                    convolution.inputs,
                    convolution.outputs,
                )
                innerproduct.add_param(
                    0, checked_cast(int, convolution.params[0].value)
                )
                innerproduct.add_param(
                    1, checked_cast(int, convolution.params[5].value)
                )
                innerproduct.add_param(
                    2, checked_cast(int, convolution.params[6].value)
                )
                innerproduct.add_param(
                    8, checked_cast(int, convolution.params[8].value)
                )
                innerproduct.add_param(
                    9, checked_cast(int, convolution.params[9].value)
                )
                innerproduct.add_param(
                    10,
                    checked_cast(list, convolution.params[10].value),
                )
                innerproduct.add_weight(
                    "weight",
                    convolution.weight_data["weight"].weight,
                    convolution.weight_data["weight"].quantize_tag,
                )
                innerproduct.add_weight("bias", convolution.weight_data["bias"].weight)

                self.model.layers[j] = innerproduct

    def __replace_convolution_with_innerproduct_after_innerproduct(self):
        while True:
            replaced = False
            for i, layer in enumerate(self.model.layers):
                if layer.op_type == "InnerProduct":
                    # InnerProduct - Convolution
                    inprod_output = layer.outputs[0]

                    j = i
                    for j in range(i + 1, len(self.model.layers)):
                        if self.model.layers[j].op_type != "Convolution":
                            continue
                        if self.model.layers[j].num_inputs != 1:
                            continue
                        if self.model.layers[j].inputs[0] == inprod_output:
                            break
                    else:
                        j += 1

                    if j == len(self.model.layers):
                        continue

                    convolution = self.model.layers[j]
                    innerproduct2 = NcnnLayer(
                        "InnerProduct",
                        convolution.name,
                        convolution.num_inputs,
                        convolution.num_outputs,
                        convolution.inputs,
                        convolution.outputs,
                    )
                    innerproduct2.add_param(
                        0, checked_cast(int, convolution.params[0].value)
                    )
                    innerproduct2.add_param(
                        1, checked_cast(int, convolution.params[5].value)
                    )
                    innerproduct2.add_param(
                        2, checked_cast(int, convolution.params[6].value)
                    )
                    innerproduct2.add_param(
                        8, checked_cast(int, convolution.params[8].value)
                    )
                    innerproduct2.add_param(
                        9, checked_cast(int, convolution.params[9].value)
                    )
                    innerproduct2.add_param(
                        10,
                        checked_cast(list, convolution.params[10].value),
                    )
                    innerproduct2.add_weight(
                        "weight",
                        convolution.weight_data["weight"].weight,
                        convolution.weight_data["weight"].quantize_tag,
                    )
                    innerproduct2.add_weight(
                        "bias", convolution.weight_data["bias"].weight
                    )

                    self.model.layers[j] = innerproduct2

                    replaced = True

            if not replaced:
                break

    def optimize(self) -> None:
        self.__fuse_batchnorm_scale()
        self.__fuse_x_batchnorm()
        self.__fuse_x_mul()
        self.__fuse_x_add()
        self.__fuse_innerproduct_dropout()

        self.__replace_reduction_with_global_pooling()
        self.__replace_prelu_with_leaky_relu()

        self.__fuse_x_activation()
        self.__fuse_memorydata_binaryop()
        self.__fuse_binaryop_eltwise()

        self.__eliminate_dropout()
        self.__eliminate_pooling1x1()
        self.__eliminate_noop()
        self.__eliminate_split()
        self.__eliminate_flatten_after_global_pooling()
        self.__eliminate_reshape_after_global_pooling()
        self.__eliminate_reshape_before_binaryop()

        self.__replace_convolution_with_innerproduct_after_global_pooling()
        self.__replace_convolution_with_innerproduct_after_innerproduct()

        self.__eliminate_flatten_after_innerproduct()
        self.__eliminate_orphaned_memorydata()

nodes/impl/ncnn/optimizer.py
vladmandic/sd-extension-chainner
from __future__ import annotations
import tempfile
from weakref import WeakKeyDictionary

try:
    from ncnn_vulkan import ncnn
    use_gpu = True
except ImportError:
    from ncnn import ncnn
    use_gpu = False


def create_ncnn_net(model, gpu_index: int) -> ncnn.Net:
    net = ncnn.Net()

    if model.fp == "fp16":
        net.opt.use_fp16_packed = True
        net.opt.use_fp16_storage = True
        net.opt.use_fp16_arithmetic = True
    else:
        net.opt.use_fp16_packed = False
        net.opt.use_fp16_storage = False
        net.opt.use_fp16_arithmetic = False

    if use_gpu:
        # Use vulkan compute
        net.opt.use_vulkan_compute = True
        net.set_vulkan_device(gpu_index)

    # Load model param and bin
    net.load_param_mem(model.model.write_param())
    if use_gpu:
        net.load_model_mem(model.model.bin)
    else:
        with tempfile.TemporaryDirectory() as tmp_model_dir:
            bin_filename = tmp_model_dir + "/ncnn-model.bin"
            model.model.write_bin(bin_filename)
            net.load_model(bin_filename)

    return net


__session_cache: WeakKeyDictionary[any, ncnn.Net] = WeakKeyDictionary()


def get_ncnn_net(model, gpu_index: int) -> ncnn.Net:
    cached = __session_cache.get(model)
    if cached is None:
        cached = create_ncnn_net(model, gpu_index)
        __session_cache[model] = cached
    return cached

nodes/impl/ncnn/session.py
vladmandic/sd-extension-chainner
from enum import Enum
from typing import Callable, List

import numpy as np

from ..utils.utils import get_h_w_c
from .image_utils import as_target_channels


def __add_noises(
    image: np.ndarray,
    noise_gen: Callable[[int, int], List[np.ndarray]],
    combine: Callable[[np.ndarray, List[np.ndarray]], np.ndarray],
) -> np.ndarray:
    img = image
    h, w, c = get_h_w_c(img)
    assert c != 2, "Noise cannot be added to 2-channel images."

    if c > 3:
        img = img[:, :, :3]

    noises = noise_gen(h, w)

    assert len(noises) > 0

    max_channels = min(c, 3)
    for n in noises:
        noise_channels = get_h_w_c(n)[2]
        assert noise_channels in (1, 3), "Noise must be a grayscale or RGB image."
        max_channels = max(max_channels, noise_channels)

    noises = [as_target_channels(n, max_channels) for n in noises]
    img = as_target_channels(img, max_channels)

    result = combine(img, noises)

    if c > 3:
        result = np.concatenate([result, image[:, :, 3:]], axis=2)

    return np.clip(result, 0, 1)


def __add_noise(
    image: np.ndarray,
    noise_gen: Callable[[int, int], np.ndarray],
    combine: Callable[[np.ndarray, np.ndarray], np.ndarray] = lambda i, n: i + n,
) -> np.ndarray:
    return __add_noises(
        image,
        lambda h, w: [noise_gen(h, w)],
        lambda i, n: combine(i, n[0]),
    )


class NoiseColor(Enum):
    RGB = "rgb"
    GRAY = "gray"

    @property
    def channels(self):
        return 3 if self is NoiseColor.RGB else 1


# Applies gaussian noise to an image
def gaussian_noise(
    image: np.ndarray,
    amount: float,
    noise_color: NoiseColor,
    seed: int = 0,
) -> np.ndarray:
    rng = np.random.default_rng(seed)
    return __add_noise(
        image,
        lambda h, w: rng.normal(0, amount, (h, w, noise_color.channels)).astype(
            np.float32
        ),
    )


# Applies uniform noise to an image
def uniform_noise(
    image: np.ndarray,
    amount: float,
    noise_color: NoiseColor,
    seed: int = 0,
) -> np.ndarray:
    rng = np.random.default_rng(seed)
    return __add_noise(
        image,
        lambda h, w: rng.uniform(-amount, amount, (h, w, noise_color.channels)).astype(
            np.float32
        ),
    )


# Applies salt and pepper noise to an image
def salt_and_pepper_noise(
    image: np.ndarray,
    amount: float,
    noise_color: NoiseColor,
    seed: int = 0,
) -> np.ndarray:
    def gen_noise(h: int, w: int):
        rng = np.random.default_rng(seed)
        noise_c = noise_color.channels
        amt = amount / 2
        pepper = rng.choice([0, 1], (h, w, noise_c), p=[amt, 1 - amt]).astype(np.uint8)
        salt = rng.choice([0, 1], (h, w, noise_c), p=[1 - amt, amt]).astype(np.uint8)
        return [pepper, salt]

    def combine(i: np.ndarray, n: List[np.ndarray]):
        pepper, salt = n
        return np.where(salt == 1, 1, np.where(pepper == 0, 0, i))

    return __add_noises(image, gen_noise, combine)


# Applies poisson noise to an image
def poisson_noise(
    image: np.ndarray,
    amount: float,
    noise_color: NoiseColor,
    seed: int = 0,
) -> np.ndarray:
    rng = np.random.default_rng(seed)
    return __add_noise(
        image,
        lambda h, w: rng.poisson(amount, (h, w, noise_color.channels)).astype(np.uint8),
    )


# Applies speckle noise to an image
def speckle_noise(
    image: np.ndarray,
    amount: float,
    noise_color: NoiseColor,
    seed: int = 0,
) -> np.ndarray:
    rng = np.random.default_rng(seed)
    return __add_noise(
        image,
        lambda h, w: rng.normal(0, amount, (h, w, noise_color.channels)).astype(
            np.float32
        ),
        lambda i, n: i + i * n,
    )

nodes/impl/noise.py
vladmandic/sd-extension-chainner
# http://momentsingraphics.de/BlueNoise.html
# https://github.com/MomentsInGraphics/BlueNoise/blob/master/BlueNoise.py

from __future__ import annotations

from typing import Tuple

import numpy as np
from scipy import ndimage


def find_largest_void(binary_pattern: np.ndarray, standard_deviation: float):
    """This function returns the indices of the largest void in the given binary
     pattern as defined by Ulichney.
    @param BinaryPattern A boolean array (should be two-dimensional although the
           implementation works in arbitrary dimensions).
    @param StandardDeviation The standard deviation used for the Gaussian filter
           in pixels. This can be a single float for an isotropic Gaussian or a
           tuple with one float per dimension for an anisotropic Gaussian.
    @return A flat index i such that BinaryPattern.flat[i] corresponds to the
            largest void. By definition this is a majority pixel.
    @sa GetVoidAndClusterBlueNoise"""
    # The minority value is always True for convenience
    if np.count_nonzero(binary_pattern) * 2 >= np.size(binary_pattern):
        binary_pattern = np.logical_not(binary_pattern)
    # Apply the Gaussian. We do not want to cut off the Gaussian at all because even
    # the tiniest difference can change the ranking. Therefore we apply the Gaussian
    # through a fast Fourier transform by means of the convolution theorem.
    filtered_array = np.fft.ifftn(
        ndimage.fourier.fourier_gaussian(
            np.fft.fftn(np.where(binary_pattern, 1.0, 0.0)), standard_deviation
        )
    ).real
    # Find the largest void
    return np.argmin(np.where(binary_pattern, 2.0, filtered_array))


def find_tightest_cluster(binary_pattern: np.ndarray, standard_deviation: float):
    """Like FindLargestVoid() but finds the tightest cluster which is a minority
     pixel by definition.
    @sa GetVoidAndClusterBlueNoise"""
    if np.count_nonzero(binary_pattern) * 2 >= np.size(binary_pattern):
        binary_pattern = np.logical_not(binary_pattern)
    filtered_array = np.fft.ifftn(
        ndimage.fourier.fourier_gaussian(
            np.fft.fftn(np.where(binary_pattern, 1.0, 0.0)), standard_deviation
        )
    ).real
    return np.argmax(np.where(binary_pattern, filtered_array, -1.0))


def create_blue_noise(
    output_shape: Tuple[int, int],
    standard_deviation: float = 1.5,
    InitialSeedFraction: float = 0.1,
    seed: int = 0,
):
    """Generates a blue noise dither array of the given shape using the method
     proposed by Ulichney [1993] in "The void-and-cluster method for dither array
     generation" published in Proc. SPIE 1913.
    @param OutputShape The shape of the output array. This function works in
           arbitrary dimension, i.e. OutputShape can have arbitrary length. Though
           it is only tested for the 2D case where you should pass a tuple
           (Height,Width).
    @param StandardDeviation The standard deviation in pixels used for the
           Gaussian filter defining largest voids and tightest clusters. Larger
           values lead to more low-frequency content but better isotropy. Small
           values lead to more ordered patterns with less low-frequency content.
           Ulichney proposes to use a value of 1.5. If you want an anisotropic
           Gaussian, you can pass a tuple of length len(OutputShape) with one
           standard deviation per dimension.
    @param InitialSeedFraction The only non-deterministic step in the algorithm
           marks a small number of pixels in the grid randomly. This parameter
           defines the fraction of such points. It has to be positive but less
           than 0.5. Very small values lead to ordered patterns, beyond that there
           is little change.
    @return An integer array of shape OutputShape containing each integer from 0
            to np.prod(OutputShape)-1 exactly once."""
    n_rank = np.prod(output_shape)
    # Generate the initial binary pattern with a prescribed number of ones
    n_initial_one = max(
        1, min(int((n_rank - 1) / 2), int(n_rank * InitialSeedFraction))
    )
    # Start from white noise (this is the only randomized step)
    initial_binary_pattern = np.zeros(output_shape, dtype=np.bool_)
    initial_binary_pattern.flat = (
        np.random.default_rng(seed).permutation(np.arange(n_rank)) < n_initial_one
    )  # type:ignore
    # Swap ones from tightest clusters to largest voids iteratively until convergence
    while True:
        i_tightest_cluster = find_tightest_cluster(
            initial_binary_pattern, standard_deviation
        )
        initial_binary_pattern.flat[i_tightest_cluster] = False
        i_largest_void = find_largest_void(initial_binary_pattern, standard_deviation)
        if i_largest_void == i_tightest_cluster:
            initial_binary_pattern.flat[i_tightest_cluster] = True
            # Nothing has changed, so we have converged
            break
        else:
            initial_binary_pattern.flat[i_largest_void] = True
    # Rank all pixels
    dither_array = np.zeros(output_shape, dtype=np.int32)
    # Phase 1: Rank minority pixels in the initial binary pattern
    binary_pattern = np.copy(initial_binary_pattern)
    for rank in range(n_initial_one - 1, -1, -1):
        i_tightest_cluster = find_tightest_cluster(binary_pattern, standard_deviation)
        binary_pattern.flat[i_tightest_cluster] = False
        dither_array.flat[i_tightest_cluster] = rank
    # Phase 2: Rank the remainder of the first half of all pixels
    binary_pattern = initial_binary_pattern
    for rank in range(n_initial_one, int((n_rank + 1) / 2)):
        i_largest_void = find_largest_void(binary_pattern, standard_deviation)
        binary_pattern.flat[i_largest_void] = True
        dither_array.flat[i_largest_void] = rank
    # Phase 3: Rank the last half of pixels
    for rank in range(int((n_rank + 1) / 2), n_rank):
        i_tightest_cluster = find_tightest_cluster(binary_pattern, standard_deviation)
        binary_pattern.flat[i_tightest_cluster] = True
        dither_array.flat[i_tightest_cluster] = rank
    return dither_array

nodes/impl/noise_functions/blue.py
vladmandic/sd-extension-chainner
"""
Simplex noise implementation by Alex Dodge, 2023

References:

Simplex noise demystified, Stefan Gustavson (2005)
http://staffwww.itn.liu.se/~stegu/simplexnoise/simplexnoise.pdf
"""

import itertools
from typing import Optional

import numpy as np

# fmt: off
PERMUTATION_TABLE_ARRAY = np.array([
    151, 160, 137, 91, 90, 15, 131, 13, 201, 95, 96, 53, 194, 233, 7, 225, 140, 36, 103,
    30, 69, 142, 8, 99, 37, 240, 21, 10, 23, 190, 6, 148, 247, 120, 234, 75, 0, 26, 197,
    62, 94, 252, 219, 203, 117, 35, 11, 32, 57, 177, 33, 88, 237, 149, 56, 87, 174, 20,
    125, 136, 171, 168, 68, 175, 74, 165, 71, 134, 139, 48, 27, 166, 77, 146, 158, 231,
    83, 111, 229, 122, 60, 211, 133, 230, 220, 105, 92, 41, 55, 46, 245, 40, 244, 102, 143,
    54, 65, 25, 63, 161, 1, 216, 80, 73, 209, 76, 132, 187, 208, 89, 18, 169, 200, 196,
    135, 130, 116, 188, 159, 86, 164, 100, 109, 198, 173, 186, 3, 64, 52, 217, 226, 250,
    124, 123, 5, 202, 38, 147, 118, 126, 255, 82, 85, 212, 207, 206, 59, 227, 47, 16, 58,
    17, 182, 189, 28, 42, 223, 183, 170, 213, 119, 248, 152, 2, 44, 154, 163, 70, 221,
    153, 101, 155, 167, 43, 172, 9, 129, 22, 39, 253, 19, 98, 108, 110, 79, 113, 224,
    232, 178, 185, 112, 104, 218, 246, 97, 228, 251, 34, 242, 193, 238, 210, 144, 12, 191,
    179, 162, 241, 81, 51, 145, 235, 249, 14, 239, 107, 49, 192, 214, 31, 181, 199, 106,
    157, 184, 84, 204, 176, 115, 121, 50, 45, 127, 4, 150, 254, 138, 236, 205, 93, 222,
    114, 67, 29, 24, 72, 243, 141, 128, 195, 78, 66, 215, 61, 156, 180
], dtype=np.int32)
# fmt: on

# Empirically determined scaling factor for different numbers of dimensions
SCALE = {
    2: 50,
    3: 39,
    4: 32,
    5: 28,
    6: 26,
    # It looks terrible at this many dimensions anyway
}


class SimplexNoise:
    def __init__(self, dimensions: int, seed: Optional[int], r2: float = 0.5):
        if dimensions <= 0:
            raise ValueError
        if dimensions == 1:
            raise RuntimeError("1D Simplex noise is not implemented here.")
        if dimensions > 6:
            raise RuntimeError("7D+ Simplex noise is not implemented here.")

        self.dimensions = dimensions
        self.r2 = r2
        self.F = (np.sqrt(self.dimensions + 1) - 1) / self.dimensions
        self.G = (1 - 1 / np.sqrt(self.dimensions + 1)) / self.dimensions

        """
        For 2D noise, we pick 16 gradients evenly distributed around the unit circle.
        For 3D and above, we pick gradients pointing at the midpoints of the edges of a hypercube centered on the origin
        """

        if self.dimensions == 2:
            n_gradients = 16
            self.gradients = np.array(
                [
                    (
                        np.cos(2 * np.pi * i / n_gradients),
                        np.sin(2 * np.pi * i / n_gradients),
                    )
                    for i in range(n_gradients)
                ]
            )
        else:
            n_gradients = self.dimensions * 2 ** (self.dimensions - 1)
            self.gradients = np.zeros((n_gradients, self.dimensions))
            for zero_dim in range(self.dimensions):
                for i, vec in enumerate(
                    itertools.product([-1, 1], repeat=self.dimensions - 1)
                ):
                    idx = zero_dim * 2 ** (self.dimensions - 1) + i
                    self.gradients[idx, :zero_dim] = vec[:zero_dim]
                    self.gradients[idx, zero_dim + 1 :] = vec[zero_dim:]

        if seed is None:
            # Use the canonical table from the reference implementation
            self.permutation_table = PERMUTATION_TABLE_ARRAY
        else:
            np.random.seed(seed)
            self.permutation_table = np.arange(self.gradients.shape[0] * 16)
            np.random.shuffle(self.permutation_table)

    def evaluate(self, points: np.ndarray):
        n_points = points.shape[0]
        assert points.shape == (n_points, self.dimensions)

        skewed_points = points + (points.sum(axis=1) * self.F).reshape((n_points, 1))
        skewed_bases, skewed_points_remainder = np.divmod(skewed_points, 1)

        skewed_simplex_verts = np.full(
            (n_points, self.dimensions + 1, self.dimensions),
            fill_value=skewed_bases.reshape((n_points, 1, -1)),
            dtype="int32",
        )

        skewed_simplex_verts[:, self.dimensions, :] += 1
        for i in range(1, self.dimensions):
            largest_dimension = np.argmax(skewed_points_remainder, axis=1)
            for o in range(self.dimensions):
                skewed_simplex_verts[
                    (largest_dimension == o), i : self.dimensions, o
                ] += 1
                if i != self.dimensions - 1:
                    skewed_points_remainder[(largest_dimension == o), o] = -1

        gradient_index = np.zeros(skewed_simplex_verts.shape[:2], dtype="int32")
        for i in range(skewed_simplex_verts.shape[2]):
            gradient_index = (
                gradient_index + skewed_simplex_verts[:, :, i]
            ) % self.permutation_table.size
            gradient_index = self.permutation_table[gradient_index]
        gradients = self.gradients[gradient_index % self.gradients.shape[0]]

        simplex_verts = (
            skewed_simplex_verts
            - skewed_simplex_verts.sum(axis=2).reshape((n_points, -1, 1)) * self.G
        )
        displacement = np.power(
            points.reshape((n_points, 1, -1)) - simplex_verts, 2
        ).sum(axis=2)
        dot_gradient = np.sum(
            (points.reshape((n_points, 1, -1)) - simplex_verts) * gradients, axis=2
        )
        contributions = (
            np.power(np.maximum(0, self.r2 - displacement), 4) * dot_gradient
        )

        return np.sum(contributions, axis=1) * SCALE[self.dimensions] + 0.5

nodes/impl/noise_functions/simplex.py
vladmandic/sd-extension-chainner
import itertools

import numpy as np


class ValueNoise:
    def __init__(self, dimensions: int, seed: int):
        self.dimensions = dimensions

        self.values = np.arange(16, dtype="float32")
        self.values = self.values / max(self.values)

        np.random.seed(seed)
        self.permutation_table = np.arange(self.values.size * 16)
        np.random.shuffle(self.permutation_table)

    def evaluate(self, points: np.ndarray):
        block, fractional = np.divmod(points, 1)

        corners = np.zeros(
            (points.shape[0], 2**self.dimensions, self.dimensions), dtype="int32"
        )
        weights = np.zeros((points.shape[0], 2**self.dimensions))

        for i, pattern in enumerate(itertools.product([0, 1], repeat=self.dimensions)):
            pattern = np.array(pattern, dtype=np.int32)
            corners[:, i] = block + pattern

            # linear interpolation
            weights[:, i] = np.prod(
                fractional * pattern + (1 - fractional) * (1 - pattern), axis=1
            )

        value_index = np.zeros(corners.shape[:2], dtype="int32")
        for i in range(corners.shape[2]):
            value_index = (value_index + corners[:, :, i]) % self.permutation_table.size
            value_index = self.permutation_table[value_index]
        values = self.values[value_index % self.values.size]

        return np.sum(values * weights, axis=1)

nodes/impl/noise_functions/value.py
vladmandic/sd-extension-chainner
import math
from enum import Enum

import numpy as np

from .util import XYZ, normalize_normals


class AdditionMethod(Enum):
    PARTIAL_DERIVATIVES = 0
    """
    The addition works by converting the normals into 2D slopes and then adding
    the slopes. The sum of the slopes is then converted back into normals.

    When adding 2 normal maps, the normals themselves are not added;
    Instead, the heightmaps that those normals represent are added.
    Conceptually, this entails converting the normals into slopes
    (the derivatives of the heightmap), integrating the slopes to get
    the heightmaps, adding the heightmaps, then performing the reverse
    on the added heightmaps. Practically, this is unnecessary, as adding
    the slopes together is equivalent to adding the heightmaps.
    """
    ANGLES = 1
    """
    The addition works by converting the normals into 2 angles, one angle the
    X axis and one along the Y axis. Those 2 angles are then added together.

    Since this might create angles outside the range of -90° to 90°,
    the resulting angles are clamped to this range.
    """


def __partial_derivatives(n1: XYZ, n2: XYZ, f1: float, f2: float) -> XYZ:
    x1, y1, z1 = n1
    x2, y2, z2 = n2

    # Slopes aren't defined for z=0, so set to near-zero decimal
    z1 = np.maximum(z1, 0.001, out=z1)
    z2 = np.maximum(z2, 0.001, out=z2)

    # This works as follows:
    # 1. Use the normals n,m to calculate 3D planes (the slopes) centered at origin p_n,p_m.
    # 2. Calculate the Z values of those planes at a_xy=(1,0) and b_xy=(0,1).
    # 3. Add the Z values to together (weighted using their strength):
    #    a_z = p_n[a_xy] * n_strength + p_m[a_xy] * m_strength, same for b_xy.
    # 4. Define a=(1,0,a_z), b=(0,1,b_z).
    # 5. The final normal will be normalize(cross(a,b)).
    # This works out as:

    n_f = f1 / z1
    m_f = f2 / z2

    x = x1 * n_f + x2 * m_f
    y = y1 * n_f + y2 * m_f

    l_r = 1 / np.sqrt(np.square(x) + np.square(y) + 1)
    x *= l_r
    y *= l_r
    z = l_r

    return x, y, z


def __clamp_angles(angles: np.ndarray) -> np.ndarray:
    return np.clip(angles, -math.pi / 2, math.pi / 2)


def __angles(n1: XYZ, n2: XYZ, f1: float, f2: float) -> XYZ:
    x1, y1, _ = n1
    x2, y2, _ = n2

    return normalize_normals(
        np.sin(__clamp_angles(np.arcsin(x1) * f1 + np.arcsin(x2) * f2)),
        np.sin(__clamp_angles(np.arcsin(y1) * f1 + np.arcsin(y2) * f2)),
    )


def add_normals(
    method: AdditionMethod,
    n1: np.ndarray,
    n2: np.ndarray,
    f1: float = 1,
    f2: float = 1,
) -> XYZ:
    # Convert BGR to XY
    x1 = n1[:, :, 2] * 2 - 1
    y1 = n1[:, :, 1] * 2 - 1
    x2 = n2[:, :, 2] * 2 - 1
    y2 = n2[:, :, 1] * 2 - 1

    xyz1 = normalize_normals(x1, y1)
    xyz2 = normalize_normals(x2, y2)

    if method is AdditionMethod.PARTIAL_DERIVATIVES:
        return __partial_derivatives(xyz1, xyz2, f1, f2)
    elif method is AdditionMethod.ANGLES:
        return __angles(xyz1, xyz2, f1, f2)
    else:
        raise AssertionError(f"Invalid normal addition method {method}")

nodes/impl/normals/addition.py
vladmandic/sd-extension-chainner
from enum import Enum
from typing import Dict, Tuple

import numpy as np


class EdgeFilter(Enum):
    SOBEL = "sobel"
    SOBEL_LIKE_5 = "sobel-like-5"
    SOBEL_LIKE_7 = "sobel-like-7"
    SOBEL_LIKE_9 = "sobel-like-9"
    PREWITT = "prewitt"
    SCHARR = "scharr"
    FOUR_SAMPLE = "4-sample"


FILTERS_X: Dict[EdgeFilter, np.ndarray] = {
    EdgeFilter.SOBEL: np.array(
        [
            [+1, 0, -1],
            [+2, 0, -2],
            [+1, 0, -1],
        ]
    ),
    EdgeFilter.SOBEL_LIKE_5: np.array(
        [
            [1 / 16, 1 / 10, 0, -1 / 10, -1 / 16],
            [1 / 10, 1 / 2.8, 0, -1 / 2.8, -1 / 10],
            [1 / 8, 1 / 2.0, 0, -1 / 2.0, -1 / 8],
            [1 / 10, 1 / 2.8, 0, -1 / 2.8, -1 / 10],
            [1 / 16, 1 / 10, 0, -1 / 10, -1 / 16],
        ]
    ),
    EdgeFilter.SOBEL_LIKE_7: np.array(
        [
            [1, 2, 3, 0, -3, -2, -1],
            [2, 3, 4, 0, -4, -3, -2],
            [3, 4, 5, 0, -5, -4, -3],
            [4, 5, 6, 0, -6, -5, -4],
            [3, 4, 5, 0, -5, -4, -3],
            [2, 3, 4, 0, -4, -3, -2],
            [1, 2, 3, 0, -3, -2, -1],
        ]
    ),
    EdgeFilter.SOBEL_LIKE_9: np.array(
        [
            [1, 2, 3, 4, 0, -4, -3, -2, -1],
            [2, 3, 4, 5, 0, -5, -4, -3, -2],
            [3, 4, 5, 6, 0, -6, -5, -4, -3],
            [4, 5, 6, 7, 0, -7, -6, -5, -4],
            [5, 6, 7, 8, 0, -8, -7, -6, -5],
            [4, 5, 6, 7, 0, -7, -6, -5, -4],
            [3, 4, 5, 6, 0, -6, -5, -4, -3],
            [2, 3, 4, 5, 0, -5, -4, -3, -2],
            [1, 2, 3, 4, 0, -4, -3, -2, -1],
        ]
    ),
    EdgeFilter.PREWITT: np.array(
        [
            [+1, 0, -1],
            [+1, 0, -1],
            [+1, 0, -1],
        ]
    ),
    EdgeFilter.SCHARR: np.array(
        [
            [+3, 0, -3],
            [+10, 0, -10],
            [+3, 0, -3],
        ]
    ),
    EdgeFilter.FOUR_SAMPLE: np.array(
        [
            [1, 0, -1],
        ]
    ),
}


def get_filter_kernels(edge_filter: EdgeFilter) -> Tuple[np.ndarray, np.ndarray]:
    """
    Returns the x and y kernels of the given edge filter.
    """

    filter_x = FILTERS_X.get(edge_filter, None)
    assert filter_x is not None, f"Unknown filter '{edge_filter}'"

    # normalize filter
    _h, w = filter_x.shape
    left = filter_x[:, : w // 2]
    filter_x = filter_x / np.sum(left)

    filter_y = np.rot90(filter_x, -1)
    return filter_x, filter_y

nodes/impl/normals/edge_filter.py
vladmandic/sd-extension-chainner
from enum import Enum

import numpy as np

from ...utils.utils import get_h_w_c


class HeightSource(Enum):
    AVERAGE_RGB = 0
    MAX_RGB = 1
    # 1 - ((1-r) * (1-g) * (1-b))
    SCREEN_RGB = 2
    RED = 3
    GREEN = 4
    BLUE = 5
    ALPHA = 6


def get_height_map(img: np.ndarray, source: HeightSource) -> np.ndarray:
    """
    Converts the given color/grayscale image to a height map.
    """
    h, w, c = get_h_w_c(img)

    assert c in (1, 3, 4), "Only grayscale, RGB, and RGBA images are supported"

    if source == HeightSource.ALPHA:
        if c < 4:
            return np.ones((h, w), dtype=np.float32)
        return img[:, :, 3]

    if c == 1:
        if source == HeightSource.SCREEN_RGB:
            x = 1 - img
            return 1 - x * x * x
        return img

    r = img[:, :, 2]
    g = img[:, :, 1]
    b = img[:, :, 0]

    if source == HeightSource.RED:
        return r
    elif source == HeightSource.GREEN:
        return g
    elif source == HeightSource.BLUE:
        return b
    elif source == HeightSource.MAX_RGB:
        return np.maximum(np.maximum(r, g), b)
    elif source == HeightSource.AVERAGE_RGB:
        return (r + g + b) / 3
    elif source == HeightSource.SCREEN_RGB:
        return 1 - ((1 - r) * (1 - g) * (1 - b))
    else:
        raise AssertionError(f"Invalid height source {source}.")

nodes/impl/normals/height.py
vladmandic/sd-extension-chainner
from typing import Tuple

import numpy as np

XYZ = Tuple[np.ndarray, np.ndarray, np.ndarray]
"""
The normalized XYZ components of a normal map. Z is guaranteed to be >= 0.
"""


def normalize_normals(x: np.ndarray, y: np.ndarray) -> XYZ:
    # The square of the length of X and Y
    l_sq = np.square(x) + np.square(y)

    # If the length of X and Y is >1, then make it 1
    l = np.sqrt(np.maximum(l_sq, 1))
    x /= l
    y /= l
    l_sq = np.minimum(l_sq, 1, out=l_sq)

    # Compute Z
    z = np.sqrt(1 - l_sq)

    return x, y, z


def gr_to_xyz(n: np.ndarray) -> XYZ:
    """
    Takes a BGR or BGRA image and converts it into XYZ normal components only by looking at the R and G channels.
    """

    x = n[:, :, 2] * 2 - 1
    y = n[:, :, 1] * 2 - 1

    return normalize_normals(x, y)


def xyz_to_bgr(xyz: XYZ) -> np.ndarray:
    """
    Converts the given XYZ components into an BGR image.
    """
    x, y, z = xyz

    r = (x + 1) * 0.5
    g = (y + 1) * 0.5
    b = z

    return np.dstack((b, g, r))


def octahedral_gr_to_xyz(n: np.ndarray) -> XYZ:
    """
    Takes a BGR or BGRA image of octahedral (RTX Remix) normals and converts it into XYZ normal components only by looking at the R and G channels.
    """
    r = n[:, :, 2] * 2 - 1
    g = n[:, :, 1] * 2 - 1

    x: np.ndarray = (r + g) / 2
    y: np.ndarray = r - x
    z: np.ndarray = 1 - np.abs(x) - np.abs(y)
    length = np.sqrt(np.square(x) + np.square(y) + np.square(z))
    x /= length
    y /= length
    z /= length  # type: ignore

    return x, y, z


def xyz_to_octahedral_bgr(xyz: XYZ) -> np.ndarray:
    """
    Converts the given XYZ components into an BGR image with normals using octahedral (RTX Remix) encoding.

    For more information about octahedral normals, see:
    https://knarkowicz.wordpress.com/2014/04/16/octahedron-normal-vector-encoding/
    https://jcgt.org/published/0003/02/01/
    """
    x, y, z = xyz
    absolute = np.abs(x) + np.abs(y) + np.abs(z)
    x /= absolute
    y /= absolute

    # This is a trick used in RTX Remix to more efficiently encode normals.
    # Octahedral normals are defined for the whole range of values (so all possible unit vectors).
    # However, we know that we are working with hemispheric normal maps (z>=0)
    # and can use this knowledge to remap xy values to assume all possible values in [-1..1]x[-1..1].
    r = x + y
    g = x - y

    r = (r + 1) * 0.5
    g = (g + 1) * 0.5
    b = np.zeros(x.shape, dtype=np.float32)

    return np.dstack((b, g, r))

nodes/impl/normals/util.py
vladmandic/sd-extension-chainner
from __future__ import annotations
import gc
import numpy as np
from ..upscale.auto_split import Tiler, auto_split
from .np_tensor_utils import np2nptensor, nptensor2np


def onnx_auto_split(
    img: np.ndarray,
    session,
    change_shape: bool,
    tiler: Tiler,
) -> np.ndarray:
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    is_fp16_model = session.get_inputs()[0].type == "tensor(float16)"

    def upscale(img: np.ndarray, _):
        try:
            lr_img = np2nptensor(img, change_range=False)
            lr_img = lr_img.astype(np.float16) if is_fp16_model else lr_img

            if change_shape:
                # Transpose from BCHW to BHWC
                lr_img = np.transpose(lr_img, (0, 2, 3, 1))

            output: np.ndarray = session.run([output_name], {input_name: lr_img})[0]

            if change_shape:
                # Transpose back to BCHW
                output = np.transpose(output, (0, 3, 1, 2))

            return nptensor2np(output, change_range=False, imtype=np.float32)
        except Exception as e:
            if "ONNXRuntimeError" in str(e) and (
                "allocate memory" in str(e)
                or "out of memory" in str(e)
                or "cudaMalloc" in str(e)
            ):
                # pylint: disable=raise-missing-from
                raise RuntimeError(
                    "A VRAM out-of-memory error has occurred. Please try using a more extreme tiling mode."
                )
            else:
                # Re-raise the exception if not an OOM error
                raise

    try:
        return auto_split(img, upscale, tiler)
    finally:
        gc.collect()

nodes/impl/onnx/auto_split.py
vladmandic/sd-extension-chainner
# This class defines an interface.
# It is important that is does not contain types that depend on ONNX.
from typing import Union

import re2

re2_options = re2.Options()
re2_options.dot_nl = True
re2_options.encoding = re2.Options.Encoding.LATIN1

U2NET_STANDARD = re2.compile(b"1959.+1960.+1961.+1962.+1963.+1964.+1965", re2_options)
U2NET_CLOTH = re2.compile(
    b"output.+d1.+Concat_1876.+Concat_1896.+Concat_1916.+Concat_1936.+Concat_1956",
    re2_options,
)
U2NET_SILUETA = re2.compile(b"1808.+1827.+1828.+2296.+1831.+1850.+1958", re2_options)
U2NET_ISNET = re2.compile(
    b"/stage1/rebnconvin/conv_s1/Conv.+/stage1/rebnconvin/relu_s1/Relu", re2_options
)


class OnnxGeneric:
    def __init__(self, model_as_bytes: bytes):
        self.bytes: bytes = model_as_bytes
        self.sub_type = "Generic"
        self.scale_height = None
        self.scale_width = None


class OnnxRemBg:
    def __init__(self, model_as_bytes: bytes, scale_height: int = 1):
        self.bytes: bytes = model_as_bytes
        self.sub_type = "RemBg"
        self.scale_height = scale_height
        self.scale_width = 1


OnnxModels = (OnnxGeneric, OnnxRemBg)
OnnxModel = Union[OnnxGeneric, OnnxRemBg]


def is_rembg_model(model_as_bytes: bytes) -> bool:
    if (
        U2NET_STANDARD.search(model_as_bytes[-600:]) is not None
        or U2NET_CLOTH.search(model_as_bytes[-1000:]) is not None
        or U2NET_SILUETA.search(model_as_bytes[-600:]) is not None
        or U2NET_ISNET.search(model_as_bytes[:10000]) is not None
    ):
        return True
    return False


def load_onnx_model(model_as_bytes: bytes) -> OnnxModel:
    if (
        U2NET_STANDARD.search(model_as_bytes[-1000:]) is not None
        or U2NET_SILUETA.search(model_as_bytes[-600:]) is not None
        or U2NET_ISNET.search(model_as_bytes[:10000]) is not None
    ):
        model = OnnxRemBg(model_as_bytes)
    elif U2NET_CLOTH.search(model_as_bytes[-1000:]) is not None:
        model = OnnxRemBg(model_as_bytes, scale_height=3)
    else:
        model = OnnxGeneric(model_as_bytes)

    return model

nodes/impl/onnx/model.py
vladmandic/sd-extension-chainner
from typing import Tuple, Type

import numpy as np

from ..image_utils import MAX_VALUES_BY_DTYPE, as_3d


def np_denorm(x: np.ndarray, min_max: Tuple[float, float] = (-1.0, 1.0)) -> np.ndarray:
    """Denormalize from [-1,1] range to [0,1]
    formula: xi' = (xi - mu)/sigma
    Example: "out = (x + 1.0) / 2.0" for denorm
        range (-1,1) to (0,1)
    for use with proper act in Generator output (ie. tanh)
    """
    out = (x - min_max[0]) / (min_max[1] - min_max[0])
    return np.clip(out, 0, 1)


def np_norm(x: np.ndarray) -> np.ndarray:
    """Normalize (z-norm) from [0,1] range to [-1,1]"""
    out = (x - 0.5) * 2.0
    return np.clip(out, -1, 1)


def np_bgr_to_rgb(img: np.ndarray) -> np.ndarray:
    out: np.ndarray = img[::-1, ...]
    return out


def np_rgb_to_bgr(img: np.ndarray) -> np.ndarray:
    # same operation as bgr_to_rgb(), flip image channels
    return np_bgr_to_rgb(img)


def np_bgra_to_rgba(img: np.ndarray) -> np.ndarray:
    out: np.ndarray = img[[2, 1, 0, 3], ...]  # type: ignore
    return out


def np_rgba_to_bgra(img: np.ndarray) -> np.ndarray:
    # same operation as bgra_to_rgba(), flip image channels
    return np_bgra_to_rgba(img)


def np2nptensor(
    img: np.ndarray,
    bgr2rgb=True,
    data_range=1.0,  # pylint: disable=unused-argument
    normalize=False,
    change_range=True,
    add_batch=True,
) -> np.ndarray:
    """Converts a numpy image array into a numpy Tensor array.
    Parameters:
        img (numpy array): the input image numpy array
        add_batch (bool): choose if new tensor needs batch dimension added
    """
    # check how many channels the image has, then condition. ie. RGB, RGBA, Gray
    # if bgr2rgb:
    #     img = img[
    #         :, :, [2, 1, 0]
    #     ]  # BGR to RGB -> in numpy, if using OpenCV, else not needed. Only if image has colors.
    if change_range:
        dtype = img.dtype
        maxval = MAX_VALUES_BY_DTYPE.get(dtype.name, 1.0)
        t_dtype = np.dtype("float32")
        img = img.astype(t_dtype) / maxval  # ie: uint8 = /255
    # "HWC to CHW" and "numpy to tensor"
    img = np.ascontiguousarray(np.transpose(as_3d(img), (2, 0, 1))).astype(np.float32)
    if bgr2rgb:
        # BGR to RGB -> in tensor, if using OpenCV, else not needed. Only if image has colors.)
        if (
            img.shape[0] % 3 == 0
        ):  # RGB or MultixRGB (3xRGB, 5xRGB, etc. For video tensors.)
            img = np_bgr_to_rgb(img)
        elif img.shape[0] == 4:  # RGBA
            img = np_bgra_to_rgba(img)
    if add_batch:
        img = np.expand_dims(
            img, axis=0
        )  # Add fake batch dimension = 1 . squeeze() will remove the dimensions of size 1
    if normalize:
        img = np_norm(img)
    return img


def nptensor2np(
    img: np.ndarray,
    rgb2bgr=True,
    remove_batch=True,
    data_range=255,
    denormalize=False,
    change_range=True,
    imtype: Type = np.uint8,
) -> np.ndarray:
    """Converts a Tensor array into a numpy image array.
    Parameters:
        img (tensor): the input image tensor array
            4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order
        remove_batch (bool): choose if tensor of shape BCHW needs to be squeezed
        denormalize (bool): Used to denormalize from [-1,1] range back to [0,1]
        imtype (type): the desired type of the converted numpy array (np.uint8
            default)
    Output:
        img (np array): 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)
    """
    n_dim = img.ndim

    img = img.astype(np.float32)

    if n_dim in (4, 3):
        # if n_dim == 4, has to convert to 3 dimensions
        if n_dim == 4 and remove_batch:
            # remove a fake batch dimension
            img = img.squeeze(0)

        if img.shape[0] == 3 and rgb2bgr:  # RGB
            # RGB to BGR -> in tensor, if using OpenCV, else not needed. Only if image has colors.
            img_np = np_rgb_to_bgr(img)
        elif img.shape[0] == 4 and rgb2bgr:  # RGBA
            # RGBA to BGRA -> in tensor, if using OpenCV, else not needed. Only if image has colors.
            img_np = np_rgba_to_bgra(img)
        else:
            img_np = img
        img_np = np.transpose(img_np, (1, 2, 0))  # CHW to HWC
    elif n_dim == 2:
        img_np = img
    else:
        raise TypeError(
            f"Only support 4D, 3D and 2D tensor. But received with dimension: {n_dim:d}"
        )

    # if rgb2bgr:
    # img_np = img_np[[2, 1, 0], :, :] #RGB to BGR -> in numpy, if using OpenCV, else not needed. Only if image has colors.
    # TODO: Check: could denormalize in the begining in tensor form instead
    if denormalize:
        img_np = np_denorm(img_np)  # denormalize if needed
    if change_range:
        img_np = np.clip(
            data_range * img_np, 0, data_range  # type: ignore
        ).round()  # np.clip to the data_range

    # has to be in range (0,255) before changing to np.uint8, else np.float32
    return img_np.astype(imtype)

nodes/impl/onnx/np_tensor_utils.py
vladmandic/sd-extension-chainner
from typing import Dict, List, Union

import numpy as np
import onnx.numpy_helper as onph
from google.protobuf.internal.containers import (
    RepeatedCompositeFieldContainer,
    RepeatedScalarFieldContainer,
)
from onnx.onnx_pb import AttributeProto, GraphProto, ModelProto, NodeProto, TensorProto
from nodes.log import logger

from ..ncnn.model import (
    DTYPE_FP16,
    DTYPE_FP32,
    BinaryOpTypes,
    EltwiseOpTypes,
    GruDirectionFlags,
    InterpResizeTypes,
    NcnnLayer,
    NcnnModel,
    NormalizeEpsModes,
    PaddingTypes,
    PadModes,
    PermuteOrderTypes,
    ReductionOpTypes,
    UnaryOpTypes,
)
from ..ncnn.optimizer import NcnnOptimizer
from .tensorproto_utils import *

UOT = UnaryOpTypes
BOT = BinaryOpTypes
EOT = EltwiseOpTypes
GRU = GruDirectionFlags
IRT = InterpResizeTypes
NEM = NormalizeEpsModes
PAM = PadModes
PAT = PaddingTypes
POT = PermuteOrderTypes
ROT = ReductionOpTypes


class Onnx2NcnnConverter:
    def __init__(self, onnx_model: ModelProto):
        self.onnx_graph: GraphProto = onnx_model.graph
        self.mutable_graph_nodes: List[NodeProto] = list(self.onnx_graph.node)
        self.node_count: int = len(self.onnx_graph.node)
        self.weights: Dict[str, TensorProto] = {
            initializer.name: initializer for initializer in self.onnx_graph.initializer
        }

        self.producers: Dict[str, None] = {i.name: None for i in self.onnx_graph.input}
        self.node_reference: Dict[str, int] = {}
        self.blob_names: Dict[str, None] = {}

    @staticmethod
    def add_weight(
        layer: NcnnLayer,
        weight_name: str,
        data: Union[float, int, np.ndarray, TensorProto],
        quantize_tag: bytes = b"",
    ) -> int:
        if isinstance(data, TensorProto):
            data = onph.to_array(data)

        return layer.add_weight(weight_name, data, quantize_tag)

    @staticmethod
    def clear_container(
        container: Union[RepeatedCompositeFieldContainer, RepeatedScalarFieldContainer],
    ) -> None:
        for _ in range(len(container)):
            container.pop()

    def swap_nodes(self, a: int, b: int) -> None:
        self.mutable_graph_nodes[a], self.mutable_graph_nodes[b] = (
            self.mutable_graph_nodes[b],
            self.mutable_graph_nodes[a],
        )

    def fuse_rewrite_gather(self) -> None:
        for gather in self.mutable_graph_nodes:
            if gather.op_type == "Gather":
                indices = get_node_attr_from_input_ai(self.weights[gather.input[1]])
                if len(indices) == 1:
                    # Reconstruct node connections
                    self.node_reference[gather.input[1]] -= 1
                    origin_inp = gather.input[0]
                    gather.ClearField("input")
                    gather.input.append(origin_inp)

                    # Update axis, starts and ends
                    axis = get_node_attr_i(gather, "axis", 1)
                    gather.op_type = "Crop"
                    gather.ClearField("attribute")

                    index = indices[0]
                    set_node_attr_ai(gather, "starts", np.array([index], np.int32))
                    set_node_attr_ai(gather, "ends", np.array([index + 1], np.int32))
                    set_node_attr_ai(gather, "axis", np.array([axis], np.int32))

    def fuse_weight_reshape(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]
            if node.op_type == "Reshape":
                if node.input[0] in self.weights:
                    self.weights[node.output[0]] = self.weights[node.input[0]]
                    if len(node.input) == 1:
                        shape = get_node_attr_ai(node, "shape")
                    elif len(node.input) == 2:
                        shape = get_node_attr_from_input_ai(self.weights[node.input[1]])
                    else:
                        shape = np.empty(0, np.int64)

                    self.clear_container(self.weights[node.output[0]].dims)
                    for dim in shape:
                        self.weights[node.output[0]].dims.append(dim)

                    node.op_type = "noop_reducedncnn"

                    self.node_reference[node.input[0]] -= 1
                    if len(node.input) == 2:
                        self.node_reference[node.input[1]] -= 1

                    reduced_node_count[0] += 1
                    i += 1

    def fuse_weight_transpose(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]
            if node.op_type == "Transpose":
                if (
                    node.input[0] in self.weights
                    and len(self.weights[node.input[0]].dims) == 2
                ):
                    perm = get_node_attr_ai(node, "perm")
                    if perm.size != 2 or perm[0] != 1 or perm[1] != 0:
                        continue

                    self.weights[node.output[0]] = self.weights[node.input[0]]

                    # Permute weight
                    B = self.weights[node.output[0]]

                    h, w = B.dims[:2]

                    permuted_data = onph.to_array(B).T

                    B.dims[:2] = (w, h)

                    if B.raw_data:
                        B.raw_data = permuted_data.tobytes()
                    else:
                        self.clear_container(B.float_data)
                        B.float_data.extend(permuted_data)

                    # Reduce
                    node.op_type = "noop_reducednccn"
                    self.node_reference[node.input[0]] -= 1

                    reduced_node_count[0] += 1
                    i += 1

    def fuse_shufflechannel(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # ShuffleChannel <= Reshape - Transpose - Reshape
            # ShuffleChannel <= Reshape - Transpose - Constant - Reshape
            if node.op_type == "Reshape":
                if self.node_reference[node.output[0]] != 1:
                    continue

                if len(node.input) == 1:
                    shape = get_node_attr_ai(node, "shape")
                else:
                    # Skip weight reshape
                    if node.input[1] not in self.weights:
                        continue
                    shape = get_node_attr_from_input_ai(self.weights[node.input[1]])

                # 1 groups channels_per_group, height, width
                # reverse style = channels_per_group, groups, height * width
                if (shape.size != 5 and shape.size != 3) or (
                    shape.size == 5 and shape[0] != 1
                ):
                    continue
                if i + 2 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]

                if node3.op_type == "Constant":
                    if i + 3 >= self.node_count:
                        continue
                    node3 = self.mutable_graph_nodes[i + 3]
                if (node2.op_type != "Transpose" or node3.op_type != "Reshape") or (
                    self.node_reference[node2.output[0]] != 1
                ):
                    continue

                # 0 2 1 3 4
                # reverse style = 1 0 2
                perm = get_node_attr_ai(node2, "perm")
                if perm.size != 5 and perm.size != 3:
                    continue
                if perm.size == 5 and (
                    perm[0] != 0
                    or perm[1] != 2
                    or perm[2] != 1
                    or perm[3] != 3
                    or perm[4] != 4
                ):
                    continue
                if perm.size == 3 and (perm[0] != 1 or perm[1] != 0 or perm[2] != 2):
                    continue

                if len(node3.input) == 1:
                    shape3 = get_node_attr_ai(node3, "shape")
                else:
                    if node3.input[1] not in self.weights:
                        continue
                    shape3 = get_node_attr_from_input_ai(self.weights[node3.input[1]])

                # 1, -1, height, width
                # reverse style = group, -1, channels_per_group, height, width
                if shape3.size != 4 and shape3.size != 5:
                    continue
                if shape3.size == 4 and (
                    shape3[0] != 1
                    or (shape3[1] != -1 and shape3[1] != shape[1] * shape[2])
                ):
                    continue
                if shape3.size == 5 and (
                    shape3[0] != shape[1]
                    or shape3[2] != shape[0]
                    or shape3[3] * shape3[4] != shape[2]
                ):
                    continue

                # Reduce
                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"

                if len(node.input) == 2:
                    self.node_reference[node.input[1]] -= 1
                self.node_reference[node.output[0]] -= 1
                self.node_reference[node2.output[0]] -= 1
                if len(node3.input) == 2:
                    self.node_reference[node3.input[1]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)

                node3.op_type = "ShuffleChannel"
                node3.input[0] = node.input[0]

                attr_group = AttributeProto(name="group", i=shape[1], type=APT.INT)
                node3.attribute.append(attr_group)

                attr_reverse = AttributeProto(
                    name="reverse", i=int(shape.size == 3), type=APT.INT
                )
                node3.attribute.append(attr_reverse)

                reduced_node_count[0] += 2
                i += 2

    def fuse_shufflechannel_split(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # Split <= ShuffleChannel(reverse type) - Gather(0) - Gather(1)
            if node.op_type == "ShuffleChannel":
                # reverse = 1
                reverse = get_node_attr_i(node, "reverse")
                if reverse != 1 or (i + 2 >= self.node_count):
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]

                if node2.op_type != "Gather" or node3.op_type != "Gather":
                    continue
                if node2.input[0] != node.output[0] or node3.input[0] != node.output[0]:
                    continue

                # axis = 0 or indices = 0
                gather2_axis = get_node_attr_i(node2, "axis")
                if gather2_axis != 0 or node2.input[1] not in self.weights:
                    continue

                gather2_indices = get_node_attr_from_input_ai(
                    self.weights[node2.input[1]]
                )
                if gather2_indices.size != 1 or gather2_indices[0] != 0:
                    continue

                # axis = 0 or indices = 1
                gather3_axis = get_node_attr_i(node3, "axis")
                if gather3_axis != 0 or node3.input[1] not in self.weights:
                    continue

                gather3_indices = get_node_attr_from_input_ai(
                    self.weights[node3.input[1]]
                )
                if gather3_indices.size != 1 or gather2_indices[0] != 1:
                    continue

                # reduce
                node2.op_type = "noop_reducedncnn"

                self.node_reference[node.output[0]] -= 2
                self.node_reference[node2.input[1]] -= 1
                self.node_reference[node3.input[1]] -= 1

                node3.op_type = "Split"
                node3.ClearField("input")
                node3.input.append(node.output[0])
                node3.output.append(node3.output[0])
                node3.output[0] = node2.output[0]

                node3.ClearField("attribute")
                attr_axis = AttributeProto(name="axis", i=1, type=APT.INT)
                node3.attribute.append(attr_axis)

                reduced_node_count[0] += 1
                i += 1

    def fuse_hardswish(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # HardSwish <= Add(+3) - Clip(0, 6) - Mul(X, ) - Div( / 6)
            # HardSwish <= Add(+3) - Clip(0, 6) - Mul(X, ) - Mul(*(1 / 6))
            # HardSwish <= Add(+3) - Clip(0, 6) - Mul(X, ) - Constant - Div( / 6)
            # HardSwish <= Add(+3) - Clip(0, 6) - Mul(X, ) - Constant - Mul(*(1 / 6))
            # out = x * F.relu6(x + 3, inplace=True) / 6
            if node.op_type == "Add":
                if (
                    self.node_reference[node.output[0]] != 1
                    or i + 3 >= self.node_count
                    or node.input[1] not in self.weights
                ):
                    continue

                add_three = self.weights[node.input[1]]
                if (
                    len(add_three.dims) != 0
                    or get_tensor_proto_data_size(add_three, add_three.data_type) != 1
                ):
                    continue

                constant_add_three = get_node_attr_from_input_f(add_three)
                if constant_add_three != 3:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]
                node4 = self.mutable_graph_nodes[i + 3]

                if node4.op_type == "Constant":
                    if i + 4 >= self.node_count:
                        continue
                    node4 = self.mutable_graph_nodes[i + 4]
                if (
                    node2.op_type != "Clip"
                    or node3.op_type != "Mul"
                    or (node4.op_type != "Div" and node4.op_type != "Mul")
                ):
                    continue
                if self.node_reference[node2.output[0]] != 1:
                    continue

                if len(node2.input) == 1:
                    relu6_min = get_node_attr_f(node2, "min", -FLOAT32_MAX)
                    relu6_max = get_node_attr_f(node2, "max", FLOAT32_MAX)
                else:
                    min_tp = self.weights[node2.input[1]]
                    max_tp = self.weights[node2.input[2]]
                    relu6_min = get_node_attr_from_input_f(min_tp)
                    relu6_max = get_node_attr_from_input_f(max_tp)

                if relu6_min != 0 or relu6_max != 6:
                    continue
                if self.node_reference[node3.output[0]] != 1:
                    continue
                if node3.input[0] != node.input[0] or node3.input[1] != node2.output[0]:
                    continue
                if node4.input[1] not in self.weights:
                    continue

                div_six = self.weights[node4.input[1]]
                if (
                    len(div_six.dims) != 0
                    or get_tensor_proto_data_size(div_six, div_six.data_type) != 1
                ):
                    continue

                constant_div_six = get_node_attr_from_input_f(div_six)
                if (node4.op_type == "Div" and constant_div_six != 6) or (
                    node4.op_type == "Mul" and constant_div_six != 1 / 6
                ):
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"
                node3.op_type = "noop_reducedncnn"

                self.node_reference[node.input[0]] -= 1
                self.node_reference[node.input[1]] -= 1
                self.node_reference[node.output[0]] -= 1
                if len(node2.input) == 3:
                    self.node_reference[node2.input[1]] -= 1
                    self.node_reference[node2.input[2]] -= 1
                self.node_reference[node2.output[0]] -= 1
                self.node_reference[node3.output[0]] -= 1
                self.node_reference[node4.input[1]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)
                self.blob_names.pop(node3.output[0], None)

                node4.op_type = "HardSwish"
                node4.ClearField("input")
                node4.input.append(node.input[0])

                attr_alpha = AttributeProto(name="alpha", f=1 / 6, type=APT.FLOAT)
                node4.attribute.append(attr_alpha)

                attr_beta = AttributeProto(name="beta", f=0.5, type=APT.FLOAT)
                node4.attribute.append(attr_beta)

                reduced_node_count[0] += 3
                i += 3

        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # HardSwish <= HardSigmoid - Mul
            # out = x * hsigmoid(x)
            if node.op_type == "HardSigmoid":
                if self.node_reference[node.output[0]] != 1:
                    continue

                alpha = get_node_attr_f(node, "alpha", 0.2)
                beta = get_node_attr_f(node, "beta", 0.5)

                if i + 1 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]

                if node2.op_type != "Mul":
                    continue
                if node2.input[0] != node.input[0] or node2.input[1] != node.output[0]:
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"

                self.node_reference[node.input[0]] -= 1
                self.node_reference[node.output[0]] -= 1

                self.blob_names.pop(node.output[0], None)

                node2.op_type = "HardSwish"
                node2.ClearField("input")
                node2.input.append(node.input[0])

                attr_alpha = AttributeProto(name="alpha", f=alpha, type=APT.FLOAT)
                node2.attribute.append(attr_alpha)

                attr_beta = AttributeProto(name="beta", f=beta, type=APT.FLOAT)
                node2.attribute.append(attr_beta)

                reduced_node_count[0] += 1
                i += 1

    def fuse_hardsigmoid(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # HardSigmoid <= Add(+3) - Clip(0, 6) - Div( / 6)
            # HardSigmoid <= Add(+3) - Clip(0, 6) - Mul(*(1 / 6))
            # HardSigmoid <= Add(+3) - Clip(0, 6) - Constant - Div( / 6)
            # HardSigmoid <= Add(+3) - Clip(0, 6) - Constant - Mul(*(1 / 6))
            # out = F.relu6(x + 3, inplace=True) / 6
            if node.op_type == "Add":
                if (
                    self.node_reference[node.output[0]] != 1
                    or i + 2 >= self.node_count
                    or node.input[1] not in self.weights
                ):
                    continue

                add_three = self.weights[node.input[1]]
                if (
                    len(add_three.dims) != 0
                    or get_tensor_proto_data_size(add_three, add_three.data_type) != 1
                ):
                    continue

                constant_add_three = self.weights[node.input[1]]
                if constant_add_three != 3:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]

                if node3.op_type == "Constant":
                    if i + 3 >= self.node_count:
                        continue
                    node3 = self.mutable_graph_nodes[i + 3]

                if node2.op_type != "Clip" or (
                    node3.op_type != "Div" and node3.op_type != "Mul"
                ):
                    continue

                if self.node_reference[node2.output[0]] != 1:
                    continue

                if len(node2.input) == 1:
                    relu6_min = get_node_attr_f(node2, "min", -FLOAT32_MAX)
                    relu6_max = get_node_attr_f(node2, "max", FLOAT32_MAX)
                else:
                    min_tp = self.weights[node2.input[1]]
                    max_tp = self.weights[node2.input[2]]
                    relu6_min = get_node_attr_from_input_f(min_tp)
                    relu6_max = get_node_attr_from_input_f(max_tp)

                if relu6_min != 0 or relu6_max != 6:
                    continue
                if node3.input[1] not in self.weights:
                    continue

                div_six = self.weights[node3.input[1]]
                if (
                    len(div_six.dims) != 0
                    or get_tensor_proto_data_size(div_six, div_six.data_type) != 1
                ):
                    continue

                constant_div_six = get_node_attr_from_input_f(div_six)
                if (node3.op_type == "Div" and constant_div_six != 6) or (
                    node3.op_type == "Mul" and constant_div_six != 1 / 6
                ):
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"

                self.node_reference[node.input[1]] -= 1
                self.node_reference[node.output[0]] -= 1
                if len(node2.input) == 3:
                    self.node_reference[node2.input[1]] -= 1
                    self.node_reference[node2.input[2]] -= 1
                self.node_reference[node2.output[0]] -= 1
                self.node_reference[node3.input[1]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)

                node3.op_type = "HardSigmoid"
                node3.ClearField("input")
                node3.input.append(node.input[0])

                attr_alpha = AttributeProto(name="alpha", f=1 / 6, type=APT.FLOAT)
                node3.attribute.append(attr_alpha)

                attr_beta = AttributeProto(name="beta", f=0.5, type=APT.FLOAT)
                node3.attribute.append(attr_beta)

                reduced_node_count[0] += 2
                i += 2

    def fuse_swish(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # Swish <= Sigmoid - Mul
            # x * torch.sigmoid(x)
            if node.op_type == "Sigmoid":
                if self.node_reference[node.output[0]] != 1 or i + 1 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]

                if node2.op_type != "Mul":
                    continue
                if node2.input[0] != node.input[0] or node2.input[1] != node.output[0]:
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"

                self.node_reference[node.input[0]] -= 1
                self.node_reference[node.output[0]] -= 1

                self.blob_names.pop(node.output[0], None)

                node2.op_type = "Swish"
                node2.ClearField("input")
                node2.input.append(node.input[0])

                reduced_node_count[0] += 1
                i += 1

    def fuse_batchnorm1d_squeeze_unsqueeze(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # BatchNormalization <= Unsqueeze - BatchNormalization - Squeeze
            if node.op_type == "Unsqueeze":
                if self.node_reference[node.output[0]] != 1 or i + 2 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]

                if node2.op_type != "BatchNormalization" or node3.op_type != "Squeeze":
                    continue
                if self.node_reference[node2.output[0]] != 1:
                    continue
                if (
                    node2.input[0] != node.output[0]
                    or node3.input[0] != node2.output[0]
                ):
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"
                node3.op_type = "noop_reducedncnn"

                self.node_reference[node.output[0]] -= 1
                self.node_reference[node2.output[0]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)

                node2.input[0] = node.input[0]
                node2.output[0] = node3.output[0]

                reduced_node_count[0] += 2
                i += 2

    def fuse_unsqueeze_prelu(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # PReLU <= Unsqueeze - PReLU
            if node.op_type == "Unsqueeze":
                # check weight
                if node.input[0] not in self.weights:
                    continue

                B = self.weights[node.input[0]]
                if len(B.dims) != 1:
                    continue
                if self.node_reference[node.output[0]] != 1:
                    continue

                # axes = (1, 2)
                axes = get_node_attr_ai(node, "axes")
                if axes.size != 2 or axes[0] != 1 or axes[1] != 2:
                    continue
                if i + 1 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]

                if node2.op_type != "PRelu" or node2.input[1] != node.output[0]:
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"

                self.node_reference[node.output[0]] -= 1

                self.blob_names.pop(node.output[0], None)

                node2.input[1] = node.input[0]

                reduced_node_count[0] += 1
                i += 1

    def fuse_normalize(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # Normalize <= X - ReduceL2 - Clip - Expand - Div
            # Normalize <= X - ReduceL2 - Clip - Shape - Expand - Div
            if node.op_type == "ReduceL2":
                if self.node_reference[node.output[0]] != 1:
                    continue

                # axes = (1)
                axes = get_node_attr_ai(node, "axes")
                if len(axes) != 1 or axes[0] != 1 or i + 3 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]
                node4 = self.mutable_graph_nodes[i + 3]

                has_shape_node = node3.op_type == "Shape"
                node_shape = NodeProto()
                if has_shape_node:
                    if i + 4 >= self.node_count:
                        continue

                    node_shape = node3
                    node3 = self.mutable_graph_nodes[i + 3]
                    node4 = self.mutable_graph_nodes[i + 4]

                if (
                    node2.op_type != "Clip"
                    or node3.op_type != "Expand"
                    or node4.op_type != "Div"
                ):
                    continue
                if (
                    self.node_reference[node2.output[0]] != 1
                    or self.node_reference[node3.output[0]] != 1
                ):
                    continue
                if (
                    node2.input[0] != node.output[0]
                    or node3.input[0] != node2.output[0]
                    or node4.input[0] != node.input[0]
                    or node4.input[1] != node3.output[0]
                ):
                    continue

                if has_shape_node and (
                    node_shape.input[0] != node.input[0]
                    or node3.input[1] != node_shape.output[0]
                ):
                    continue

                # +eps
                if len(node2.input) == 1:
                    clip_min = get_node_attr_f(node2, "min", -FLOAT32_MAX)
                else:
                    min_tp = self.weights[node2.input[1]]
                    clip_min = get_node_attr_from_input_f(min_tp)

                # reduce
                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"
                if has_shape_node:
                    node_shape.op_type = "noop_reducedncnn"
                node3.op_type = "noop_reducedncnn"

                self.node_reference[node.input[0]] -= 2 if has_shape_node else 1
                self.node_reference[node.output[0]] -= 1
                self.node_reference[node2.output[0]] -= 1
                if has_shape_node:
                    self.node_reference[node_shape.output[0]] -= 1
                self.node_reference[node3.output[0]] -= 1
                if len(node3.input) == 2:
                    self.node_reference[node3.input[1]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)
                if has_shape_node:
                    self.blob_names.pop(node_shape.output[0], None)
                self.blob_names.pop(node3.output[0], None)

                node4.op_type = "Normalize"
                node4.ClearField("input")
                node4.input.append(node.input[0])

                attr_alpha = AttributeProto(name="eps", f=clip_min, type=APT.FLOAT)
                node4.attribute.append(attr_alpha)

                reduced_node_count[0] += 4 if has_shape_node else 3
                i += 4 if has_shape_node else 3

    def fuse_groupnorm(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # GroupNorm <= X - Reshape - InstanceNormalization - Reshape - Mul - Add
            if node.op_type == "Reshape":
                if self.node_reference[node.output[0]] != 1:
                    continue

                if len(node.input) == 1:
                    shape = get_node_attr_ai(node, "shape")
                else:
                    # Skip weight reshape
                    if node.input[1] not in self.weights:
                        continue

                    shape = get_node_attr_from_input_ai(self.weights[node.input[1]])

                # 0, group, -1
                if (
                    shape.size != 3
                    or shape[0] != 0
                    or shape[2] != -1
                    or i + 4 >= self.node_count
                ):
                    continue

                groups = shape[1]

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]
                node4 = self.mutable_graph_nodes[i + 3]
                node5 = self.mutable_graph_nodes[i + 4]

                if (
                    node2.op_type != "InstanceNormalization"
                    or node3.op_type != "Reshape"
                    or node4.op_type != "Mul"
                    or node5.op_type != "Add"
                ):
                    continue
                if (
                    self.node_reference[node2.output[0]] != 1
                    or self.node_reference[node3.output[0]] != 1
                    or self.node_reference[node4.output[0]] != 1
                ):
                    continue
                if (
                    node2.input[0] != node.output[0]
                    or node3.input[0] != node2.output[0]
                    or node4.input[0] != node3.output[0]
                    or node5.input[0] != node4.output[0]
                ):
                    continue

                # InstanceNormalization S=1 B=0
                S = get_node_attr_from_input_af(self.weights[node2.input[1]])
                B = get_node_attr_from_input_af(self.weights[node2.input[2]])
                if S.size != groups or B.size != groups:
                    continue
                if np.any(S != 1) or np.any(B != 0):
                    continue

                if len(node3.input) == 1:
                    shape2 = get_node_attr_ai(node3, "shape")
                else:
                    # Skip weight reshape
                    if node3.input[1] not in self.weights:
                        continue

                    shape2 = get_node_attr_from_input_ai(self.weights[node3.input[1]])

                # 1, channels, w, h
                if shape2.size != 4 or shape2[0] != 1:
                    continue

                channels = shape2[1]

                # affine
                affine_S = get_node_attr_from_input_af(self.weights[node4.input[1]])
                affine_B = get_node_attr_from_input_af(self.weights[node5.input[1]])
                if affine_S.size != channels and affine_B.size != channels:
                    continue  # only per-channel affine allowed

                # reduce
                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"
                node3.op_type = "noop_reducedncnn"
                node4.op_type = "noop_reducedncnn"

                if len(node.input) == 2:
                    self.node_reference[node.input[1]] -= 1
                self.node_reference[node.output[0]] -= 1
                self.node_reference[node2.input[1]] -= 1
                self.node_reference[node2.input[2]] -= 1
                self.node_reference[node2.output[0]] -= 1
                if len(node3.input) == 2:
                    self.node_reference[node3.input[1]] -= 1
                self.node_reference[node3.output[0]] -= 1
                self.node_reference[node4.output[0]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)
                self.blob_names.pop(node3.output[0], None)
                self.blob_names.pop(node4.output[0], None)

                affine_scale = node4.input[1]
                affine_bias = node5.input[1]

                node5.op_type = "GroupNorm"
                node5.ClearField("input")
                node5.input.append(node.input[0])
                node5.input.append(affine_scale)
                node5.input.append(affine_bias)

                attr_groups = AttributeProto(name="groups", i=groups, type=APT.INT)
                node5.attribute.append(attr_groups)

                attr_channels = AttributeProto(
                    name="channels", i=channels, type=APT.INT
                )
                node5.attribute.append(attr_channels)

                # +eps
                eps = get_node_attr_f(node2, "epsilon", 0.00001)
                attr_eps = AttributeProto(name="epsilon", f=eps, type=APT.FLOAT)
                node5.attribute.append(attr_eps)

                attr_affine = AttributeProto(name="affine", i=1, type=APT.INT)
                node5.attribute.append(attr_affine)

                reduced_node_count[0] += 4
                i += 4

    def fuse_layernorm(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # LayerNorm <= X - ReduceMean - Sub - Pow - ReduceMean - Add - Sqrt - Div
            # LayerNorm <= X - ReduceMean - Sub - Pow - ReduceMean - Add - Sqrt - Div - Mul - Add
            if node.op_type == "ReduceMean":
                if self.node_reference[node.output[0]] != 1:
                    continue

                axes = get_node_attr_ai(node, "axes")

                # -1
                # -2 -1
                if axes.size != 1 and axes.size != 2:
                    continue
                if (axes.size == 1 and axes[0] != -1) or (
                    axes.size == 2 and (axes[0] != -2 or axes[1] != -1)
                ):
                    continue
                if i + 6 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]
                node4 = self.mutable_graph_nodes[i + 3]
                node5 = self.mutable_graph_nodes[i + 4]
                node6 = self.mutable_graph_nodes[i + 5]
                node7 = self.mutable_graph_nodes[i + 6]

                if node2.op_type != "Sub" or node3.op_type != "Pow":
                    continue
                if (
                    self.node_reference[node2.output[0]] != 2
                    or self.node_reference[node3.output[0]] != 1
                    or self.node_reference[node4.output[0]] != 1
                    or self.node_reference[node5.output[0]] != 1
                    or self.node_reference[node6.output[0]] != 1
                ):
                    continue
                if (
                    node2.input[0] != node.output[0]
                    or node2.input[1] != node.output[0]
                    or node3.input[0] != node2.output[0]
                    or node4.input[0] != node3.output[0]
                    or node5.input[0] != node4.output[0]
                    or node6.input[0] != node5.output[0]
                    or node7.input[0] != node2.output[0]
                    or node7.input[1] != node6.output[0]
                ):
                    continue
                if node3.input[1] not in self.weights:
                    continue

                pow_two = self.weights[node3.input[1]]
                if (
                    len(pow_two.dims) != 0
                    or get_tensor_proto_data_size(pow_two, pow_two.data_type) != 1
                ):
                    continue

                constant_pow_two = get_node_attr_from_input_f(pow_two)
                if constant_pow_two != 2:
                    continue

                axes4 = get_node_attr_ai(node4, "axes")

                # -1
                # -2 -1
                if axes4.size != axes.size:
                    continue
                if (axes.size == 1 and axes[4] != -1) or (
                    axes.size == 2 and (axes4[0] != -2 or axes4[1] != -1)
                ):
                    continue
                if node5.input[1] not in self.weights:
                    continue

                add_eps = self.weights[node5.input[1]]
                if (
                    len(add_eps.dims) != 0
                    or get_tensor_proto_data_size(add_eps, add_eps.data_type) != 1
                ):
                    continue

                eps = get_node_attr_from_input_f(add_eps)

                affine = 0
                while i + 8 < self.node_count:
                    node8 = self.mutable_graph_nodes[i + 7]
                    node9 = self.mutable_graph_nodes[i + 8]

                    if node8.op_type != "Mul" or node9.op_type != "Add":
                        break
                    if (
                        self.node_reference[node7.output[0]] != 1
                        or self.node_reference[node8.output[0]] != 1
                    ):
                        break
                    if (
                        node8.input[0] != node7.output[0]
                        or node9.input[0] != node8.output[0]
                    ):
                        break

                    # affine
                    affine_S = get_node_attr_from_input_af(self.weights[node8.input[1]])
                    affine_B = get_node_attr_from_input_af(self.weights[node9.input[1]])
                    if affine_S.size != affine_B.size:
                        break

                    affine = 1
                    break

                # reduce
                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"
                node3.op_type = "noop_reducedncnn"
                node4.op_type = "noop_reducedncnn"
                node5.op_type = "noop_reducedncnn"
                node6.op_type = "noop_reducedncnn"

                self.node_reference[node2.input[0]] -= 1
                self.node_reference[node2.input[1]] -= 1
                self.node_reference[node3.input[0]] -= 1
                self.node_reference[node3.input[1]] -= 1
                self.node_reference[node4.input[0]] -= 1
                self.node_reference[node5.input[0]] -= 1
                self.node_reference[node5.input[1]] -= 1
                self.node_reference[node6.input[0]] -= 1
                self.node_reference[node7.input[0]] -= 1
                self.node_reference[node7.input[1]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)
                self.blob_names.pop(node3.output[0], None)
                self.blob_names.pop(node4.output[0], None)
                self.blob_names.pop(node5.output[0], None)
                self.blob_names.pop(node6.output[0], None)

                attr_eps = AttributeProto(name="epsilon", f=eps, type=APT.FLOAT)
                attr_affine = AttributeProto(name="affine", i=affine, type=APT.INT)
                if affine == 0:
                    node7.op_type = "LayerNorm"
                    node7.ClearField("input")
                    node7.input.append(node.input[0])

                    node7.attribute.append(attr_eps)
                    node7.attribute.append(attr_affine)

                    reduced_node_count[0] += 6
                    i += 6
                else:
                    # This is probably unnecessary on their part, but I'm paranoid
                    node8 = self.mutable_graph_nodes[i + 7]
                    node9 = self.mutable_graph_nodes[i + 8]

                    node7.op_type = "noop_reducedncnn"
                    node8.op_type = "noop_reducedncnn"

                    self.node_reference[node8.input[0]] -= 1
                    self.node_reference[node9.input[0]] -= 1

                    self.blob_names.pop(node7.output[0], None)
                    self.blob_names.pop(node8.output[0], None)

                    affine_scale = node8.input[1]
                    affine_bias = node9.input[1]

                    node9.op_type = "LayerNorm"
                    node9.ClearField("input")
                    node9.input.append(node.input[0])
                    node9.input.append(affine_scale)
                    node9.input.append(affine_bias)

                    node9.attribute.append(attr_eps)
                    node9.attribute.append(attr_affine)

                    reduced_node_count[0] += 8
                    i += 8

    def fuse_flatten(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # Flatten <= X - Shape - Gather - Constant - Unsqueeze - Unsqueeze - Concat - Reshape
            if node.op_type == "Shape":
                if self.node_reference[node.output[0]] != 1:
                    continue
                if i + 6 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]
                node4 = self.mutable_graph_nodes[i + 3]
                node5 = self.mutable_graph_nodes[i + 4]
                node6 = self.mutable_graph_nodes[i + 5]
                node7 = self.mutable_graph_nodes[i + 6]

                if (
                    node2.op_type != "Gather"
                    or node3.op_type != "Constant"
                    or node4.op_type != "Unsqueeze"
                    or node5.op_type != "Unsqueeze"
                    or node6.op_type != "Concat"
                    or node7.op_type != "Reshape"
                ):
                    continue
                if (
                    self.node_reference[node2.output[0]] != 1
                    or self.node_reference[node4.output[0]] != 1
                    or self.node_reference[node5.output[0]] != 1
                    or self.node_reference[node6.output[0]] != 1
                ):
                    continue
                if (
                    node2.input[0] != node.output[0]
                    or node4.input[0] != node2.output[0]
                    or node5.input[0] != node3.output[0]
                    or node6.input[0] != node4.output[0]
                    or node6.input[1] != node5.output[0]
                    or node7.input[0] != node.input[0]
                    or node7.input[1] != node6.output[0]
                ):
                    continue

                # axis = 0
                gather_axis = get_node_attr_i(node2, "axis")
                if gather_axis != 0:
                    continue

                # indices = 0
                if node2.input[1] not in self.weights:
                    continue

                gather_indices = get_node_attr_from_input_ai(
                    self.weights[node2.input[1]]
                )
                if gather_indices.size != 1 or gather_indices[0] != 0:
                    continue

                # axes = (0)
                unsqueeze_axes = get_node_attr_ai(node4, "axes")
                if unsqueeze_axes.size != 1 or unsqueeze_axes[0] != 0:
                    continue
                unsqueeze_axes2 = get_node_attr_ai(node5, "axes")
                if unsqueeze_axes2.size != 1 or unsqueeze_axes2[0] != 0:
                    continue

                # data = -1
                if node5.input[0] not in self.weights:
                    continue

                unsqueeze2_data = get_node_attr_from_input_ai(
                    self.weights[node5.input[0]]
                )
                if unsqueeze2_data.size != 1 or unsqueeze2_data[0] != -1:
                    continue

                # axis = 0
                concat_axis = get_node_attr_i(node6, "axis")
                if concat_axis != 0:
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"
                node4.op_type = "noop_reducedncnn"
                node5.op_type = "noop_reducedncnn"
                node6.op_type = "noop_reducedncnn"

                self.node_reference[node.input[0]] -= 1
                self.node_reference[node.output[0]] -= 1
                self.node_reference[node2.input[1]] -= 1
                self.node_reference[node2.output[0]] -= 1
                self.node_reference[node4.output[0]] -= 1
                self.node_reference[node5.input[0]] -= 1
                self.node_reference[node5.output[0]] -= 1
                self.node_reference[node.output[0]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)
                self.blob_names.pop(node4.output[0], None)
                self.blob_names.pop(node5.output[0], None)
                self.blob_names.pop(node6.output[0], None)

                node7.op_type = "Flatten"
                node7.ClearField("input")
                node7.input.append(node.input[0])

                reduced_node_count[0] += 5
                i += 5

    def fuse_pixelshuffle(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # PixelShuffle <= Reshape - Transpose - Reshape
            # PixelShuffle <= Reshape - Transpose - Constant - Reshape
            if node.op_type == "Reshape":
                if self.node_reference[node.output[0]] != 1:
                    continue

                if len(node.input) == 1:
                    shape = get_node_attr_ai(node, "shape")
                else:
                    # skip weight reshape
                    if node.input[1] not in self.weights:
                        continue

                    shape = get_node_attr_from_input_ai(self.weights[node.input[1]])

                # -1, 3, upscale_factor, upscale_factor, height, width
                if (
                    shape.size != 6
                    or (shape[0] != 1 and shape[0] != -1)
                    or shape[2] != shape[3]
                    or i + 2 >= self.node_count
                ):
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]

                if node3.op_type == "Constant":
                    if i + 3 >= self.node_count:
                        continue

                    node3 = self.mutable_graph_nodes[i + 3]

                if node2.op_type != "Transpose" or node3.op_type != "Reshape":
                    continue
                if self.node_reference[node2.output[0]] != 1:
                    continue

                # 0 1 4 2 5 3
                perm = get_node_attr_ai(node2, "perm")
                if (
                    perm.size != 6
                    or perm[0] != 0
                    or perm[1] != 1
                    or perm[2] != 4
                    or perm[3] != 2
                    or perm[4] != 5
                    or perm[5] != 3
                ):
                    continue

                if len(node3.input) == 1:
                    shape3 = get_node_attr_ai(node3, "shape")
                else:
                    if node3.input[1] not in self.weights:
                        continue

                    shape3 = get_node_attr_from_input_ai(self.weights[node3.input[1]])

                # -1, 3, height, width
                if (
                    shape3.size != 4
                    or (shape3[0] != 1 and shape3[0] != -1)
                    or shape3[1] != shape[1]
                    or shape3[2] != shape[2] * shape[4]
                    or shape3[3] != shape[3] * shape[5]
                ):
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"

                if len(node.input) == 2:
                    self.node_reference[node.input[1]] -= 1
                self.node_reference[node.output[0]] -= 1
                self.node_reference[node2.output[0]] -= 1
                if len(node3.input) == 2:
                    self.node_reference[node3.input[1]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)

                node3.op_type = "PixelShuffle"
                node3.input[0] = node.input[0]

                attr_group = AttributeProto(
                    name="scale_factor", i=shape[2], type=APT.INT
                )
                node3.attribute.append(attr_group)

                reduced_node_count[0] += 2
                i += 2

    def fuse_reorg(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # PixelShuffle <= Reshape - Transpose - Reshape
            # PixelShuffle <= Reshape - Transpose - Constant - Reshape
            if node.op_type == "Reshape":
                if self.node_reference[node.output[0]] != 1:
                    continue

                if len(node.input) == 1:
                    shape = get_node_attr_ai(node, "shape")
                else:
                    if node.input[1] not in self.weights:
                        continue

                    shape = get_node_attr_from_input_ai(self.weights[node.input[1]])

                # -1, 3, out_height, block_size, out_width, block_size
                if (
                    shape.size != 6
                    or (shape[0] != 1 and shape[0] != -1)
                    or shape[3] != shape[5]
                    or i + 2 >= self.node_count
                ):
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]

                if node3.op_type == "Constant":
                    if i + 3 >= self.node_count:
                        continue

                    node3 = self.mutable_graph_nodes[i + 3]

                if node2.op_type != "Transpose" or node3.op_type != "Reshape":
                    continue
                if self.node_reference[node2.output[0]] != 1:
                    continue

                # 0 1 3 5 2 4
                perm = get_node_attr_ai(node2, "perm")
                if (
                    perm.size != 6
                    or perm[0] != 0
                    or perm[1] != 1
                    or perm[2] != 3
                    or perm[3] != 5
                    or perm[4] != 2
                    or perm[5] != 4
                ):
                    continue

                if len(node3.input) == 1:
                    shape3 = get_node_attr_ai(node3, "shape")
                else:
                    if node3.input[1] not in self.weights:
                        continue

                    shape3 = get_node_attr_from_input_ai(self.weights[node3.input[1]])

                # -1, out_channels, out_height, out_width
                if (
                    shape3.size != 4
                    or (shape3[0] != 1 and shape3[0] != -1)
                    or shape3[1] != shape[1] * shape[3] * shape[5]
                    or shape3[2] != shape[2]
                    or shape3[3] != shape[4]
                ):
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"

                if len(node.input) == 2:
                    self.node_reference[node.input[1]] -= 1
                self.node_reference[node.output[0]] -= 1
                self.node_reference[node2.output[0]] -= 1
                if len(node3.input) == 2:
                    self.node_reference[node3.input[1]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)

                node3.op_type = "Reorg"
                node3.input[0] = node.input[0]

                attr_group = AttributeProto(name="stride", i=shape[3], type=APT.INT)
                node3.attribute.append(attr_group)

                reduced_node_count[0] += 2
                i += 2

    def fuse_expand_broadcast(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # Add/Sub/Mul/Div/Min/Max <= Expand - Add/Sub/Mul/Div/Min/Max
            if node.op_type == "Expand":
                if self.node_reference[node.output[0]] != 1 or i + 1 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]

                if node2.op_type not in ["Add", "Sub", "Mul", "Div", "Min", "Max"]:
                    continue
                if (
                    node2.input[1] != node.output[0]
                    and node2.input[0] != node.output[0]
                ):
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"

                self.node_reference[node.output[0]] -= 1
                if len(node.input) == 2:
                    self.node_reference[node.input[1]] -= 1

                self.blob_names.pop(node.output[0], None)

                if node2.input[0] == node.output[0]:
                    node2.input[0] = node.input[0]
                else:
                    node2.input[1] = node.input[0]

                reduced_node_count[0] += 1
                i += 1

    def fuse_lstm_gru_rnn(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # LSTM(bi) <= LSTM(bi) - Transpose - Reshape - Transpose
            if node.op_type in ["LSTM", "GRU", "RNN"]:
                if self.node_reference[node.output[0]] != 1 or i + 2 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]

                if node2.op_type != "Transpose" or node3.op_type != "Reshape":
                    continue
                if self.node_reference[node2.output[0]] != 1:
                    continue
                if (
                    node2.input[0] != node.output[0]
                    or node3.input[0] != node2.output[0]
                ):
                    continue

                direction = get_node_attr_s(node, "direction")
                if direction != "bidirectional":
                    continue

                # 0 2 1 3
                perm = get_node_attr_ai(node2, "perm")
                if (
                    perm.size != 4
                    or perm[0] != 0
                    or perm[1] != 2
                    or perm[2] != 1
                    or perm[3] != 3
                ):
                    continue

                if len(node3.input) == 1:
                    shape = get_node_attr_ai(node3, "shape")
                else:
                    if node3.input[1] not in self.weights:
                        continue

                    shape = get_node_attr_from_input_ai(self.weights[node3.input[1]])

                # 0 0 -1
                if shape.size != 3 or shape[0] != 0 or shape[1] != 0 or shape[2] != -1:
                    continue

                # reduce
                node2.op_type = "noop_reducedncnn"
                node3.op_type = "noop_reducedncnn"

                self.node_reference[node.output[0]] -= 1
                self.node_reference[node2.output[0]] -= 1
                if len(node3.input) == 2:
                    self.node_reference[node3.input[1]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)

                node.output[0] = node3.output[0]

                reduced_node_count[0] += 2
                i += 2

                if i + 1 < self.node_count:
                    if self.node_reference[node3.output[0]] != 1:
                        continue

                    node4 = self.mutable_graph_nodes[i + 1]

                    if node4.op_type != "Transpose":
                        continue
                    if node4.input[0] != node.output[0]:
                        continue

                    # 1 0 2
                    perm4 = get_node_attr_ai(node4, "perm")
                    if (
                        perm4.size != 3
                        or perm4[0] != 1
                        or perm4[1] != 0
                        or perm4[2] != 2
                    ):
                        continue

                    # reduce
                    node4.op_type = "noop_reducedncnn"

                    self.node_reference[node.output[0]] -= 1

                    self.blob_names.pop(node.output[0], None)

                    node.output[0] = node4.output[0]

                    reduced_node_count[0] += 1
                    i += 1

        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # LSTM(uni) <= LSTM(uni) - Squeeze - Transpose
            if node.op_type in ["LSTM", "GRU", "RNN"]:
                if self.node_reference[node.output[0]] != 1 or i + 1 >= self.node_count:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]

                if node2.op_type != "Squeeze":
                    continue
                if node2.input[0] != node.output[0]:
                    continue

                direction = get_node_attr_s(node, "direction")
                if direction == "bidirectional":
                    continue

                axes = get_node_attr_ai(node2, "axes")
                if axes.size != 1 or axes[0] != 1:
                    continue

                # reduce
                node2.op_type = "noop_reducedncnn"

                self.node_reference[node.output[0]] -= 1

                self.blob_names.pop(node.output[0], None)

                node.output[0] = node2.output[0]

                reduced_node_count[0] += 1
                i += 1

                if i + 1 < self.node_count:
                    if self.node_reference[node2.output[0]] != 1:
                        continue

                    node3 = self.mutable_graph_nodes[i + 1]

                    if node3.op_type != "Transpose":
                        continue

                    if node3.input[0] != node.output[0]:
                        continue

                    # 1 0 2
                    perm4 = get_node_attr_ai(node3, "perm")
                    if (
                        perm4.size != 3
                        or perm4[0] != 1
                        or perm4[1] != 0
                        or perm4[2] != 2
                    ):
                        continue

                    # reduce
                    node3.op_type = "noop_reducedncnn"

                    self.node_reference[node.output[0]] -= 1

                    self.blob_names.pop(node.output[0], None)

                    node.output[0] = node3.output[0]

                    reduced_node_count[0] += 1
                    i += 1

        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # LSTM <= Transpose - LSTM
            if node.op_type == "Transpose":
                if self.node_reference[node.output[0]] != 1:
                    continue

                # 1 0 2
                perm = get_node_attr_ai(node, "perm")
                if perm.size != 3 or perm[0] != 1 or perm[1] != 0 or perm[2] != 2:
                    continue

                node2 = self.mutable_graph_nodes[i + 1]

                if node2.op_type not in ["LSTM", "GRU", "RNN"]:
                    continue
                if node2.input[0] != node.output[0]:
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"

                self.node_reference[node.output[0]] -= 1

                self.blob_names.pop(node.output[0], None)

                node2.input[0] = node.input[0]

                reduced_node_count[0] += 1
                i += 1

    def fuse_multiheadattention(self, reduced_node_count: List[int]) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # MultiHeadAttention <= MatMul(q) - Add
            #                      - MatMul(k) - Add
            #                      - MatMul(v) - Add
            #                      - Mul
            #                      - Reshape - Transpose
            #                      - Reshape - Reshape - Transpose - Transpose
            #                      - Gemm - Softmax - Gemm - Transpose - Reshape - MatMul - Add
            if node.op_type == "MatMul":
                if (
                    self.node_reference[node.output[0]] != 1
                    or i + 19 >= self.node_count
                ):
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]
                node4 = self.mutable_graph_nodes[i + 3]
                node5 = self.mutable_graph_nodes[i + 4]
                node6 = self.mutable_graph_nodes[i + 5]
                node7 = self.mutable_graph_nodes[i + 6]
                node8 = self.mutable_graph_nodes[i + 7]
                node9 = self.mutable_graph_nodes[i + 8]
                node10 = self.mutable_graph_nodes[i + 9]
                node11 = self.mutable_graph_nodes[i + 10]
                node12 = self.mutable_graph_nodes[i + 11]
                node13 = self.mutable_graph_nodes[i + 12]
                node14 = self.mutable_graph_nodes[i + 13]
                node15 = self.mutable_graph_nodes[i + 14]
                node16 = self.mutable_graph_nodes[i + 15]
                node17 = self.mutable_graph_nodes[i + 16]
                node18 = self.mutable_graph_nodes[i + 17]
                node19 = self.mutable_graph_nodes[i + 18]
                node20 = self.mutable_graph_nodes[i + 19]

                if (
                    node2.op_type != "Add"
                    or node3.op_type != "MatMul"
                    or node4.op_type != "Add"
                    or node5.op_type != "MatMul"
                    or node6.op_type != "Add"
                    or node7.op_type != "Mul"
                    or node8.op_type != "Reshape"
                    or node9.op_type != "Transpose"
                    or node10.op_type != "Reshape"
                    or node11.op_type != "Reshape"
                    or node12.op_type != "Transpose"
                    or node13.op_type != "Transpose"
                    or node14.op_type != "MatMul"
                    or node15.op_type != "Softmax"
                    or node16.op_type != "MatMul"
                    or node17.op_type != "Transpose"
                    or node18.op_type != "Reshape"
                    or node19.op_type != "MatMul"
                    or node20.op_type != "Add"
                ):
                    continue
                if (
                    self.node_reference[node2.output[0]] != 1
                    or self.node_reference[node3.output[0]] != 1
                    or self.node_reference[node4.output[0]] != 1
                    or self.node_reference[node5.output[0]] != 1
                    or self.node_reference[node6.output[0]] != 1
                    or self.node_reference[node7.output[0]] != 1
                    or self.node_reference[node8.output[0]] != 1
                    or self.node_reference[node9.output[0]] != 1
                    or self.node_reference[node10.output[0]] != 1
                    or self.node_reference[node11.output[0]] != 1
                    or self.node_reference[node12.output[0]] != 1
                    or self.node_reference[node13.output[0]] != 1
                    or self.node_reference[node14.output[0]] != 1
                    or self.node_reference[node15.output[0]] != 1
                    or self.node_reference[node16.output[0]] != 1
                    or self.node_reference[node17.output[0]] != 1
                    or self.node_reference[node18.output[0]] != 1
                    or self.node_reference[node19.output[0]] != 1
                ):
                    continue
                if (
                    node2.input[0] != node.output[0]
                    or node4.input[0] != node3.output[0]
                    or node6.input[0] != node5.output[0]
                    or node7.input[0] != node2.output[0]
                    or node8.input[0] != node7.output[0]
                    or node9.input[0] != node8.output[0]
                    or node10.input[0] != node4.output[0]
                    or node11.input[0] != node6.output[0]
                    or node12.input[0] != node11.output[0]
                    or node13.input[0] != node10.output[0]
                    or node14.input[0] != node9.output[0]
                    or node14.input[1] != node13.output[0]
                    or node15.input[0] != node14.output[0]
                    or node16.input[0] != node15.output[0]
                    or node16.input[1] != node12.output[0]
                    or node17.input[0] != node16.output[0]
                    or node18.input[0] != node17.output[0]
                    or node19.input[0] != node18.output[0]
                    or node20.input[0] != node19.output[0]
                ):
                    continue

                q_B = get_node_attr_from_input_af(self.weights[node2.input[1]])
                k_B = get_node_attr_from_input_af(self.weights[node4.input[1]])
                v_B = get_node_attr_from_input_af(self.weights[node6.input[1]])
                o_B = get_node_attr_from_input_af(self.weights[node20.input[1]])

                if q_B.size != k_B.size or q_B.size != v_B.size or q_B.size != o_B.size:
                    continue

                embed_dim = q_B.size

                # 1 0 2
                perm9 = get_node_attr_ai(node9, "perm")
                perm12 = get_node_attr_ai(node12, "perm")
                if perm9.size != 3 or perm9[0] != 1 or perm9[1] != 0 or perm9[2] != 2:
                    continue
                if (
                    perm12.size != 3
                    or perm12[0] != 1
                    or perm12[1] != 0
                    or perm12[2] != 2
                ):
                    continue

                # 1 2 0
                perm13 = get_node_attr_ai(node13, "perm")
                if (
                    perm13.size != 3
                    or perm13[0] != 1
                    or perm13[1] != 2
                    or perm13[2] != 0
                ):
                    continue

                # 1 0 2
                perm17 = get_node_attr_ai(node17, "perm")
                if (
                    perm17.size != 3
                    or perm17[0] != 1
                    or perm17[1] != 0
                    or perm17[2] != 2
                ):
                    continue

                softmax_axis = get_node_attr_i(node15, "axis")
                if softmax_axis != 2:
                    continue

                # 1/-1 seqlen * num_heads, embed_dim / num_heads
                if len(node8.input) == 1:
                    shape8 = get_node_attr_ai(node8, "shape")
                else:
                    if node8.input[1] not in self.weights:
                        continue

                    shape8 = get_node_attr_from_input_ai(self.weights[node8.input[1]])
                if len(node10.input) == 1:
                    shape10 = get_node_attr_ai(node10, "shape")
                else:
                    if node10.input[1] not in self.weights:
                        continue

                    shape10 = get_node_attr_from_input_ai(self.weights[node10.input[1]])
                if len(node11.input) == 1:
                    shape11 = get_node_attr_ai(node11, "shape")
                else:
                    if node11.input[1] not in self.weights:
                        continue

                    shape11 = get_node_attr_from_input_ai(self.weights[node11.input[1]])

                if shape8.size != 3 or shape10.size != 3 or shape11.size != 3:
                    continue
                if (
                    shape8[1] != shape10[1]
                    or shape8[1] != shape11[1]
                    or shape8[2] != shape10[2]
                    or shape8[2] != shape11[2]
                ):
                    continue

                num_heads = embed_dim / shape8[2]

                if len(node18.input) == 1:
                    shape18 = get_node_attr_ai(node18, "shape")
                else:
                    if node18.input[1] not in self.weights:
                        continue

                    shape18 = get_node_attr_from_input_ai(self.weights[node18.input[1]])

                if (
                    shape18.size != 3
                    or shape18[2] != embed_dim
                    or shape18[1] * num_heads != shape8[1]
                ):
                    continue

                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"
                node3.op_type = "noop_reducedncnn"
                node4.op_type = "noop_reducedncnn"
                node5.op_type = "noop_reducedncnn"
                node6.op_type = "noop_reducedncnn"
                node7.op_type = "noop_reducedncnn"
                node8.op_type = "noop_reducedncnn"
                node9.op_type = "noop_reducedncnn"
                node10.op_type = "noop_reducedncnn"
                node11.op_type = "noop_reducedncnn"
                node12.op_type = "noop_reducedncnn"
                node13.op_type = "noop_reducedncnn"
                node14.op_type = "noop_reducedncnn"
                node15.op_type = "noop_reducedncnn"
                node16.op_type = "noop_reducedncnn"
                node17.op_type = "noop_reducedncnn"
                node18.op_type = "noop_reducedncnn"
                node19.op_type = "noop_reducedncnn"

                self.node_reference[node2.input[0]] -= 1
                self.node_reference[node4.input[0]] -= 1
                self.node_reference[node6.input[0]] -= 1
                self.node_reference[node7.input[0]] -= 1
                self.node_reference[node7.input[1]] -= 1
                self.node_reference[node8.input[0]] -= 1
                if len(node8.input) == 2:
                    self.node_reference[node8.input[1]] -= 1
                self.node_reference[node9.input[0]] -= 1
                self.node_reference[node10.input[0]] -= 1
                if len(node10.input) == 2:
                    self.node_reference[node10.input[1]] -= 1
                self.node_reference[node11.input[0]] -= 1
                if len(node11.input) == 2:
                    self.node_reference[node11.input[1]] -= 1
                self.node_reference[node12.input[0]] -= 1
                self.node_reference[node13.input[0]] -= 1
                self.node_reference[node14.input[0]] -= 1
                self.node_reference[node14.input[1]] -= 1
                self.node_reference[node15.input[0]] -= 1
                self.node_reference[node16.input[0]] -= 1
                self.node_reference[node16.input[1]] -= 1
                self.node_reference[node17.input[0]] -= 1
                self.node_reference[node18.input[0]] -= 1
                if len(node18.input) == 2:
                    self.node_reference[node18.input[1]] -= 1
                self.node_reference[node19.input[0]] -= 1
                self.node_reference[node20.input[0]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)
                self.blob_names.pop(node3.output[0], None)
                self.blob_names.pop(node4.output[0], None)
                self.blob_names.pop(node5.output[0], None)
                self.blob_names.pop(node6.output[0], None)
                self.blob_names.pop(node7.output[0], None)
                self.blob_names.pop(node8.output[0], None)
                self.blob_names.pop(node9.output[0], None)
                self.blob_names.pop(node10.output[0], None)
                self.blob_names.pop(node11.output[0], None)
                self.blob_names.pop(node12.output[0], None)
                self.blob_names.pop(node13.output[0], None)
                self.blob_names.pop(node14.output[0], None)
                self.blob_names.pop(node15.output[0], None)
                self.blob_names.pop(node16.output[0], None)
                self.blob_names.pop(node17.output[0], None)
                self.blob_names.pop(node18.output[0], None)
                self.blob_names.pop(node19.output[0], None)

                qw = node.input[1]
                qb = node2.input[1]
                kw = node3.input[1]
                kb = node4.input[1]
                vw = node5.input[1]
                vb = node6.input[1]
                ow = node19.input[1]
                ob = node20.input[1]

                node20.op_type = "MultiHeadAttention"
                node20.ClearField("input")
                node20.input.append(node.input[0])
                node20.input.append(node3.input[0])
                node20.input.append(node5.input[0])
                node20.input.append(qw)
                node20.input.append(qb)
                node20.input.append(kw)
                node20.input.append(kb)
                node20.input.append(vw)
                node20.input.append(vb)
                node20.input.append(ow)
                node20.input.append(ob)

                attr_embed_dim = AttributeProto(
                    name="embed_dim", i=embed_dim, type=APT.INT
                )
                node20.attribute.append(attr_embed_dim)

                attr_num_heads = AttributeProto(
                    name="num_heads", i=num_heads, type=APT.INT
                )
                node20.attribute.append(attr_num_heads)

                reduced_node_count[0] += 19
                i += 19

        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # MultiHeadAttention <= MatMul(qkv) - Add - Split
            #                      - Mul
            #                      - Reshape - Transpose
            #                      - Reshape - Reshape - Transpose - Transpose
            #                      - Gemm - Softmax - Gemm - Transpose - Reshape - MatMul - Add
            if node.op_type == "MatMul":
                if (
                    self.node_reference[node.output[0]] != 1
                    or i + 16 >= self.node_count
                ):
                    continue

                node2 = self.mutable_graph_nodes[i + 1]
                node3 = self.mutable_graph_nodes[i + 2]
                node4 = self.mutable_graph_nodes[i + 3]
                node5 = self.mutable_graph_nodes[i + 4]
                node6 = self.mutable_graph_nodes[i + 5]
                node7 = self.mutable_graph_nodes[i + 6]
                node8 = self.mutable_graph_nodes[i + 7]
                node9 = self.mutable_graph_nodes[i + 8]
                node10 = self.mutable_graph_nodes[i + 9]
                node11 = self.mutable_graph_nodes[i + 10]
                node12 = self.mutable_graph_nodes[i + 11]
                node13 = self.mutable_graph_nodes[i + 12]
                node14 = self.mutable_graph_nodes[i + 13]
                node15 = self.mutable_graph_nodes[i + 14]
                node16 = self.mutable_graph_nodes[i + 15]
                node17 = self.mutable_graph_nodes[i + 16]

                if (
                    node2.op_type != "Add"
                    or node3.op_type != "Split"
                    or node4.op_type != "Mul"
                    or node5.op_type != "Reshape"
                    or node6.op_type != "Transpose"
                    or node7.op_type != "Reshape"
                    or node8.op_type != "Reshape"
                    or node9.op_type != "Transpose"
                    or node10.op_type != "Transpose"
                    or node11.op_type != "MatMul"
                    or node12.op_type != "Softmax"
                    or node13.op_type != "MatMul"
                    or node14.op_type != "Transpose"
                    or node15.op_type != "Reshape"
                    or node16.op_type != "MatMul"
                    or node17.op_type != "Add"
                ):
                    continue
                if (
                    self.node_reference[node2.output[0]] != 1
                    or self.node_reference[node3.output[0]] != 1
                    or self.node_reference[node3.output[1]] != 1
                    or self.node_reference[node3.output[2]] != 1
                    or self.node_reference[node4.output[0]] != 1
                    or self.node_reference[node5.output[0]] != 1
                    or self.node_reference[node6.output[0]] != 1
                    or self.node_reference[node7.output[0]] != 1
                    or self.node_reference[node8.output[0]] != 1
                    or self.node_reference[node9.output[0]] != 1
                    or self.node_reference[node10.output[0]] != 1
                    or self.node_reference[node11.output[0]] != 1
                    or self.node_reference[node12.output[0]] != 1
                    or self.node_reference[node13.output[0]] != 1
                    or self.node_reference[node14.output[0]] != 1
                    or self.node_reference[node15.output[0]] != 1
                    or self.node_reference[node16.output[0]] != 1
                ):
                    continue
                if (
                    node2.input[0] != node.output[0]
                    or node3.input[0] != node2.output[0]
                    or node4.input[0] != node3.output[0]
                    or node5.input[0] != node4.output[0]
                    or node6.input[0] != node5.output[0]
                    or node7.input[0] != node3.output[1]
                    or node8.input[0] != node3.output[2]
                    or node9.input[0] != node8.output[0]
                    or node10.input[0] != node7.output[0]
                    or node11.input[0] != node6.output[0]
                    or node11.input[1] != node10.output[0]
                    or node12.input[0] != node11.output[0]
                    or node13.input[0] != node12.output[0]
                    or node13.input[1] != node9.output[0]
                    or node14.input[0] != node13.output[0]
                    or node15.input[0] != node14.output[0]
                    or node16.input[0] != node15.output[0]
                    or node17.input[0] != node16.output[0]
                ):
                    continue

                qkv_B = get_node_attr_from_input_af(self.weights[node2.input[1]])
                o_B = get_node_attr_from_input_af(self.weights[node17.input[1]])

                if qkv_B.size != o_B.size * 3:
                    continue

                embed_dim = o_B.size

                # 1 0 2
                perm6 = get_node_attr_ai(node6, "perm")
                perm9 = get_node_attr_ai(node9, "perm")
                if perm6.size != 3 or perm6[0] != 1 or perm6[1] != 0 or perm6[2] != 2:
                    continue
                if perm9.size != 3 or perm9[0] != 1 or perm9[1] != 0 or perm9[2] != 2:
                    continue

                # 1 2 0
                perm10 = get_node_attr_ai(node10, "perm")
                if (
                    perm10.size != 3
                    or perm10[0] != 1
                    or perm10[1] != 2
                    or perm10[2] != 0
                ):
                    continue

                # 1 0 2
                perm14 = get_node_attr_ai(node14, "perm")
                if (
                    perm14.size != 3
                    or perm14[0] != 1
                    or perm14[1] != 0
                    or perm14[2] != 2
                ):
                    continue

                softmax_axis = get_node_attr_i(node12, "axis")
                if softmax_axis != 2:
                    continue

                # 1/-1, seqlen * num_heads, embed_dim / num_heads
                if len(node5.input) == 1:
                    shape5 = get_node_attr_ai(node5, "shape")
                else:
                    if node5.input[1] not in self.weights:
                        continue

                    shape5 = get_node_attr_from_input_ai(self.weights[node5.input[1]])
                if len(node7.input) == 1:
                    shape7 = get_node_attr_ai(node7, "shape")
                else:
                    if node7.input[1] not in self.weights:
                        continue

                    shape7 = get_node_attr_from_input_ai(self.weights[node7.input[1]])
                if len(node8.input) == 1:
                    shape8 = get_node_attr_ai(node8, "shape")
                else:
                    if node8.input[1] not in self.weights:
                        continue

                    shape8 = get_node_attr_from_input_ai(self.weights[node8.input[1]])

                if (
                    shape5[1] != shape7[1]
                    or shape5[1] != shape8[1]
                    or shape5[2] != shape7[2]
                    or shape5[2] != shape8[2]
                ):
                    continue

                num_heads = embed_dim / shape5[2]

                # 1, seqlen, embed_dim
                if len(node15.input) == 1:
                    shape15 = get_node_attr_ai(node15, "shape")
                else:
                    if node15.input[1] not in self.weights:
                        continue

                    shape15 = get_node_attr_from_input_ai(self.weights[node15.input[1]])

                if (
                    shape15.size != 3
                    or shape15[2] != embed_dim
                    or shape15[1] * num_heads != shape8[1]
                ):
                    continue

                # reduce
                node.op_type = "noop_reducedncnn"
                node2.op_type = "noop_reducedncnn"
                node3.op_type = "noop_reducedncnn"
                node4.op_type = "noop_reducedncnn"
                node5.op_type = "noop_reducedncnn"
                node6.op_type = "noop_reducedncnn"
                node7.op_type = "noop_reducedncnn"
                node8.op_type = "noop_reducedncnn"
                node9.op_type = "noop_reducedncnn"
                node10.op_type = "noop_reducedncnn"
                node11.op_type = "noop_reducedncnn"
                node12.op_type = "noop_reducedncnn"
                node13.op_type = "noop_reducedncnn"
                node14.op_type = "noop_reducedncnn"
                node15.op_type = "noop_reducedncnn"
                node16.op_type = "noop_reducedncnn"

                self.node_reference[node2.input[0]] -= 1
                self.node_reference[node3.input[0]] -= 1
                self.node_reference[node4.input[0]] -= 1
                self.node_reference[node4.input[1]] -= 1
                self.node_reference[node5.input[0]] -= 1
                if len(node5.input) == 2:
                    self.node_reference[node5.input[1]] -= 1
                self.node_reference[node6.input[0]] -= 1
                self.node_reference[node7.input[0]] -= 1
                if len(node7.input) == 2:
                    self.node_reference[node7.input[1]] -= 1
                self.node_reference[node8.input[0]] -= 1
                if len(node8.input) == 2:
                    self.node_reference[node8.input[1]] -= 1
                self.node_reference[node9.input[0]] -= 1
                self.node_reference[node10.input[0]] -= 1
                self.node_reference[node11.input[0]] -= 1
                self.node_reference[node11.input[1]] -= 1
                self.node_reference[node12.input[0]] -= 1
                self.node_reference[node13.input[0]] -= 1
                self.node_reference[node13.input[1]] -= 1
                self.node_reference[node14.input[0]] -= 1
                self.node_reference[node15.input[0]] -= 1
                if len(node15.input) == 2:
                    self.node_reference[node15.input[1]] -= 1
                self.node_reference[node16.input[0]] -= 1
                self.node_reference[node17.input[0]] -= 1

                self.blob_names.pop(node.output[0], None)
                self.blob_names.pop(node2.output[0], None)
                self.blob_names.pop(node3.output[0], None)
                self.blob_names.pop(node3.output[1], None)
                self.blob_names.pop(node3.output[2], None)
                self.blob_names.pop(node4.output[0], None)
                self.blob_names.pop(node5.output[0], None)
                self.blob_names.pop(node6.output[0], None)
                self.blob_names.pop(node7.output[0], None)
                self.blob_names.pop(node8.output[0], None)
                self.blob_names.pop(node9.output[0], None)
                self.blob_names.pop(node10.output[0], None)
                self.blob_names.pop(node11.output[0], None)
                self.blob_names.pop(node12.output[0], None)
                self.blob_names.pop(node13.output[0], None)
                self.blob_names.pop(node14.output[0], None)
                self.blob_names.pop(node15.output[0], None)
                self.blob_names.pop(node16.output[0], None)

                qkvw = node.input[1]
                qkvb = node2.input[1]
                ow = node16.input[1]
                ob = node17.input[1]

                node17.op_type = "MultiHeadAttention"
                node17.ClearField("input")
                node17.input.append(node.input[0])
                node17.input.append(qkvw)
                node17.input.append(qkvb)
                node17.input.append(ow)
                node17.input.append(ob)

                attr_embed_dim = AttributeProto(
                    name="embed_dim", i=embed_dim, type=APT.INT
                )
                node17.attribute.append(attr_embed_dim)

                attr_num_heads = AttributeProto(
                    name="num_heads", i=num_heads, type=APT.INT
                )
                node17.attribute.append(attr_num_heads)

                reduced_node_count[0] += 16
                i += 16

    def fuse_binaryop_with_scalar(self) -> None:
        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # Add/Sub/Mul/Div/Min/Max/Pow(a, x)
            if node.op_type in ["Add", "Sub", "Mul", "Div", "Min", "Max", "Pow"]:
                if node.input[0] not in self.weights:
                    continue

                scalar_b = self.weights[node.input[0]]
                if (
                    len(scalar_b.dims) != 0
                    or get_tensor_proto_data_size(scalar_b, scalar_b.data_type) != 1
                ):
                    continue

                if node.op_type == "Sub":
                    node.op_type = "RSub"
                elif node.op_type == "Div":
                    node.op_type = "RDiv"

                b = get_node_attr_from_input_f(scalar_b)

                self.node_reference[node.input[0]] -= 1

                node_input = node.input[1]
                node.ClearField("input")
                node.input.append(node_input)

                attr_with_scalar = AttributeProto(name="with_scalar", i=1, type=APT.INT)
                node.attribute.append(attr_with_scalar)

                attr_b = AttributeProto(name="b", f=b, type=APT.FLOAT)
                node.attribute.append(attr_b)

        for i in range(self.node_count):
            node = self.mutable_graph_nodes[i]

            # Add/Sub/Mul/Div/Min/Max/Pow(x, b)
            if node.op_type in ["Add", "Sub", "Mul", "Div", "Min", "Max", "Pow"]:
                if node.input[1] not in self.weights:
                    continue

                scalar_b = self.weights[node.input[1]]
                if (
                    len(scalar_b.dims) != 0
                    or get_tensor_proto_data_size(scalar_b, scalar_b.data_type) != 1
                ):
                    continue

                b = get_node_attr_from_input_f(scalar_b)

                self.node_reference[node.input[1]] -= 1

                node_input = node.input[0]
                node.ClearField("input")
                node.input.append(node_input)

                attr_with_scalar = AttributeProto(name="with_scalar", i=1, type=APT.INT)
                node.attribute.append(attr_with_scalar)

                attr_b = AttributeProto(name="b", f=b, type=APT.FLOAT)
                node.attribute.append(attr_b)

    def convert(self, is_fp16: bool = False, include_mem_data: bool = True):
        if is_fp16:
            logger.debug("NCNN mode: fp16")
        else:
            logger.debug("NCNN mode: fp32")

        # Topological sort
        i = 0
        while i < self.node_count:
            node = self.mutable_graph_nodes[i]
            swapnode = False
            missing_input_name = None
            for input_name in node.input:
                if (
                    input_name
                    and input_name not in self.producers
                    and input_name not in self.weights
                ):
                    swapnode = True
                    missing_input_name = input_name
                    break

            # If nothing missing, add outputs to producers and continue
            # to next node
            if not swapnode:
                for output_name in node.output:
                    if output_name:
                        self.producers[output_name] = None

                i += 1
                continue

            # find node that produces missing_input_name
            for j, nodeq in enumerate(self.mutable_graph_nodes, i + 1):
                found = False
                for output_name in nodeq.output:
                    if output_name == missing_input_name:
                        found = True
                        break

                if found:
                    break
            else:
                raise RuntimeError(
                    f"Cannot find node that produces {missing_input_name}, "
                    f"which is required by node {i} ({node.name})."
                )

            self.swap_nodes(i, j)

        # global definition line
        # [layer count][blob count]
        for node in self.onnx_graph.node:
            op = node.op_type
            if not node.name:
                node.name = node.output[0]

            if op == "Constant":
                self.weights[node.output[0]] = get_node_attr_tensor(node, "value")

            for input_name in node.input:
                self.blob_names[input_name] = None

                if input_name not in self.node_reference:
                    self.node_reference[input_name] = 1
                else:
                    self.node_reference[input_name] += 1

            if op == "Dropout":
                output_name = node.output[0]
                self.blob_names[output_name] = None
                self.node_reference[output_name] = 0
                continue

            for output_name in node.output:
                self.blob_names[output_name] = None
                self.node_reference[output_name] = 0

        # include Input node
        input_node_count = 0
        for graph_input in self.onnx_graph.input:
            input_name = graph_input.name

            # check weight
            if input_name not in self.weights:
                self.blob_names[input_name] = None
                input_node_count += 1

        # op chain fusion
        reduced_node_count = [0]
        self.fuse_weight_reshape(reduced_node_count)
        self.fuse_weight_transpose(reduced_node_count)
        self.fuse_shufflechannel(reduced_node_count)
        self.fuse_shufflechannel_split(reduced_node_count)
        self.fuse_hardsigmoid(reduced_node_count)
        self.fuse_hardswish(reduced_node_count)
        self.fuse_swish(reduced_node_count)
        self.fuse_batchnorm1d_squeeze_unsqueeze(reduced_node_count)
        self.fuse_unsqueeze_prelu(reduced_node_count)
        self.fuse_normalize(reduced_node_count)
        self.fuse_groupnorm(reduced_node_count)
        self.fuse_layernorm(reduced_node_count)
        self.fuse_flatten(reduced_node_count)
        self.fuse_pixelshuffle(reduced_node_count)
        self.fuse_reorg(reduced_node_count)
        self.fuse_expand_broadcast(reduced_node_count)
        self.fuse_lstm_gru_rnn(reduced_node_count)
        self.fuse_multiheadattention(reduced_node_count)
        self.fuse_binaryop_with_scalar()
        self.fuse_rewrite_gather()

        # reduce common const weight node_reference
        for node in self.onnx_graph.node:
            op = node.op_type
            if op == "BatchNormalization":
                self.node_reference[node.input[1]] -= 1
                self.node_reference[node.input[2]] -= 1
                self.node_reference[node.input[3]] -= 1
                self.node_reference[node.input[4]] -= 1
            elif op == "BiasGelu":
                self.node_reference[node.input[1]] -= 1
            elif op == "Clip":
                if len(node.input) == 3:
                    self.node_reference[node.input[1]] -= 1
                    self.node_reference[node.input[2]] -= 1
            elif op == "Conv":
                self.node_reference[node.input[1]] -= 1
                if len(node.input) == 3:
                    self.node_reference[node.input[2]] -= 1
            elif op == "ConvTranspose":
                self.node_reference[node.input[1]] -= 1
                if len(node.input) == 3:
                    self.node_reference[node.input[2]] -= 1
            elif op == "EmbedLayerNormalization":
                self.node_reference[node.input[1]] -= 1
                self.node_reference[node.input[2]] -= 1
                self.node_reference[node.input[3]] -= 1
                self.node_reference[node.input[4]] -= 1
                self.node_reference[node.input[5]] -= 1
                self.node_reference[node.input[6]] -= 1
            elif op == "Gemm":
                alpha = get_node_attr_f(node, "alpha", 1)
                beta = get_node_attr_f(node, "beta", 1)
                transA = get_node_attr_i(node, "transA", 0)
                transB = get_node_attr_i(node, "transB", 0)

                if alpha == 1 and beta == 1 and transA == 0 and transB == 1:
                    # InnerProduct-like A * B + C
                    self.node_reference[node.input[1]] -= 1
                    self.node_reference[node.input[2]] -= 1
            elif op == "GroupNorm":
                affine = get_node_attr_i(node, "affine", 1)
                if affine:
                    self.node_reference[node.input[1]] -= 1
                    self.node_reference[node.input[2]] -= 1
            elif op == "GRU":
                for gru_input in node.input:
                    self.node_reference[gru_input] -= 1
            elif op == "InstanceNormalization":
                self.node_reference[node.input[1]] -= 1
                self.node_reference[node.input[2]] -= 1
            elif op == "LayerNorm":
                affine = get_node_attr_i(node, "affine", 1)
                if affine:
                    self.node_reference[node.input[1]] -= 1
                    self.node_reference[node.input[2]] -= 1
            elif op == "LSTM":
                for lstm_input in node.input:
                    self.node_reference[lstm_input] -= 1
            elif op == "MatMul":
                if (
                    node.input[1] in self.weights
                    and len(self.weights[node.input[1]].dims) == 2
                ):
                    # InnerProduct
                    self.node_reference[node.input[1]] -= 1
            elif op == "MultiHeadAttention":
                if len(node.input) == 5:
                    self.node_reference[node.input[1]] -= 1
                    self.node_reference[node.input[2]] -= 1
                    self.node_reference[node.input[3]] -= 1
                    self.node_reference[node.input[4]] -= 1
                else:
                    self.node_reference[node.input[3]] -= 1
                    self.node_reference[node.input[4]] -= 1
                    self.node_reference[node.input[5]] -= 1
                    self.node_reference[node.input[6]] -= 1
                    self.node_reference[node.input[7]] -= 1
                    self.node_reference[node.input[8]] -= 1
                    self.node_reference[node.input[9]] -= 1
                    self.node_reference[node.input[10]] -= 1
            elif op == "Pad":
                if len(node.input) >= 2:
                    self.node_reference[node.input[1]] -= 1
            elif op == "PRelu":
                self.node_reference[node.input[1]] -= 1
            elif op == "Reshape":
                if len(node.input) >= 2:
                    self.node_reference[node.input[1]] -= 1
            elif op == "Resize":
                if len(node.input) == 2:
                    # opset 10
                    self.node_reference[node.input[1]] -= 1
                else:
                    # opset 11+
                    self.node_reference[node.input[1]] -= 1
                    self.node_reference[node.input[2]] -= 1
                    if len(node.input) >= 4:
                        self.node_reference[node.input[3]] -= 1
            elif op == "RNN":
                for rnn_input in node.input:
                    self.node_reference[rnn_input] -= 1
            elif op == "SkipLayerNormalization":
                self.node_reference[node.input[2]] -= 1
                self.node_reference[node.input[3]] -= 1
                self.node_reference[node.input[4]] -= 1
            elif op == "Slice":
                if len(node.input) >= 2:
                    self.node_reference[node.input[1]] -= 1
                    self.node_reference[node.input[2]] -= 1
                    if len(node.input) >= 4:
                        self.node_reference[node.input[3]] -= 1
                    if len(node.input) >= 5:
                        self.node_reference[node.input[4]] -= 1
            elif op == "Upsample":
                if len(node.input) >= 2:
                    self.node_reference[node.input[1]] -= 1
            elif op == "adaptive_avg_pool2d" or op == "adaptive_max_pool2d":
                if len(node.input) >= 2:
                    self.node_reference[node.input[1]] -= 1

        # count all weight node with zero reference
        zero_reference_weight_node_count = 0
        for input_name in self.weights.keys():
            # there may be some weight nodes in initializer but none of the graph nodes use them
            # add them to blob_names so we could get proper blob count later
            self.blob_names[input_name] = None

            refcount = self.node_reference[input_name]
            if refcount == 0:
                zero_reference_weight_node_count += 1

        # we always treat constant nodes as weights or binaryop_weights
        # do not count it twice for layer_count
        constant_node_count_moved_to_weight = 0
        for node in self.onnx_graph.node:
            if node.op_type == "Constant":
                constant_node_count_moved_to_weight += 1

        # some ops may have anonymous input
        # LSTM sequence_lens
        self.blob_names.pop("", None)
        self.node_reference.pop("", None)

        # remove node_reference entries with references equal to one
        split_layer_count = 0
        splitncnn_blob_count = 0

        # split node reference
        split_node_reference = {}
        for ref, count in self.node_reference.items():
            if count > 1:
                split_layer_count += 1
                splitncnn_blob_count += count
                split_node_reference[ref] = count

        ncnn_node_count = (
            self.node_count
            - constant_node_count_moved_to_weight
            + len(self.weights)
            - zero_reference_weight_node_count
            - reduced_node_count[0]
            + input_node_count
            + split_layer_count
        )
        ncnn_blob_count = (
            len(self.blob_names)
            - zero_reference_weight_node_count
            + splitncnn_blob_count
        )
        ncnn_model = NcnnModel(ncnn_node_count, ncnn_blob_count)
        logger.debug(
            f"Node count: {ncnn_model.node_count}, Blob count: {ncnn_model.blob_count}"
        )

        bin_length = 0
        for i, graph_input in enumerate(self.onnx_graph.input):
            input_name = graph_input.name

            # Make sure input is not in weights
            if input_name not in self.weights:
                ncnn_model.add_layer(
                    NcnnLayer("Input", input_name, 0, 1, outputs=[input_name])
                )

                refcount = self.node_reference[input_name]
                if refcount > 1:
                    layer_input_list = [
                        f"{input_name}_splitncnn_{j}" for j in range(refcount)
                    ]
                    ncnn_model.add_layer(
                        NcnnLayer(
                            "Split",
                            f"splitncnn_input{i}",
                            1,
                            refcount,
                            [input_name],
                            layer_input_list,
                        )
                    )

        # place MemoryData next if it is being included
        internal_split = 0
        if include_mem_data:
            for input_name, M in self.weights.items():
                refcount = self.node_reference[input_name]
                if refcount != 0:
                    layer = NcnnLayer("MemoryData", input_name, 0, 1, [input_name])

                    M_dims_size = len(M.dims)
                    if M_dims_size == 0:
                        layer.add_param(0, get_tensor_proto_data_size(M, M.data_type))
                    elif M_dims_size == 1:
                        layer.add_param(0, M.dims[0])
                    elif M_dims_size == 2:
                        layer.add_param(0, M.dims[1])
                        if M.dims[0] != 1:
                            layer.add_param(1, M.dims[0])
                    elif M_dims_size == 3:
                        layer.add_param(0, M.dims[2])
                        layer.add_param(1, M.dims[1])
                        if M.dims[0] != 1:
                            layer.add_param(2, M.dims[0])
                    elif M_dims_size == 4:
                        layer.add_param(0, M.dims[3])
                        layer.add_param(1, M.dims[2])
                        layer.add_param(2, M.dims[1])

                    bin_length += self.add_weight(layer, "MemoryData", M)

                    ncnn_model.add_layer(layer)

                    if refcount > 1:
                        layer_output_list = [
                            f"{input_name}_splitncnn_{i}" for i in range(refcount)
                        ]
                        ncnn_model.add_layer(
                            NcnnLayer(
                                "Split",
                                f"splitncnn_{internal_split}",
                                1,
                                refcount,
                                [input_name],
                                layer_output_list,
                            )
                        )

                        internal_split += 1

        for node in self.onnx_graph.node:
            op = node.op_type

            if op == "noop_reducedncnn":
                continue

            name = node.name
            if not name:
                name = node.output[0]

            input_size = len(node.input)
            output_size = len(node.output)

            for input_name in node.input:
                # check weight
                if not input_name or (
                    input_name in self.weights and self.node_reference[input_name] == 0
                ):
                    input_size -= 1

            layer = NcnnLayer()
            if op in [
                "Abs",
                "Acos",
                "Asin",
                "Atan",
                "Ceil",
                "Cos",
                "Exp",
                "Floor",
                "Log",
                "Neg",
                "Reciprocal",
                "Sin",
                "Sqrt",
                "Tan",
                "Tanh",
            ]:
                layer.op_type = "UnaryOp"
            elif op in [
                "Add",
                "Div",
                "Max",
                "Min",
                "Mul",
                "Pow",
                "RDiv",
                "RSub",
                "Sub",
            ]:
                layer.op_type = "BinaryOp"
            elif op == "AveragePool" or op == "MaxPool":
                kernel_shape = get_node_attr_ai(node, "kernel_shape")
                if kernel_shape.size == 1:
                    layer.op_type = "Pooling1D"
                else:
                    layer.op_type = "Pooling"
            elif op == "BatchNormalization":
                layer.op_type = "BatchNorm"
            elif op == "BiasGelu":
                layer.op_type = "BiasGelu"
            elif op == "Clip":
                layer.op_type = "Clip"
            elif op == "Concat":
                layer.op_type = "Concat"
            elif op == "Constant":
                continue
            elif op == "Conv":
                kernel_shape = get_node_attr_ai(node, "kernel_shape")
                if kernel_shape.size == 1:
                    layer.op_type = "Convolution1D"
                else:
                    group = get_node_attr_i(node, "group", 1)
                    if group > 1:
                        layer.op_type = "ConvolutionDepthWise"
                    else:
                        layer.op_type = "Convolution"
            elif op == "ConvTranspose":
                group = get_node_attr_i(node, "group", 1)
                if group > 1:
                    layer.op_type = "DeconvolutionDepthWise"
                else:
                    layer.op_type = "Deconvolution"
            elif op == "Crop" or op == "Slice":
                layer.op_type = "Crop"
            elif op == "DepthToSpace" or op == "PixelShuffle":
                layer.op_type = "PixelShuffle"
            elif op == "Dropout":
                layer.op_type = "Dropout"
                output_size = 1
            elif op == "Elu":
                layer.op_type = "ELU"
            elif op == "EmbedLayerNormalization":
                layer.op_type = "EmbedLayerNormalization"
            elif op == "Flatten":
                layer.op_type = "Flatten"
            elif op == "Gelu":
                layer.op_type = "GELU"
            elif op == "Gemm":
                alpha = get_node_attr_f(node, "alpha", 1)
                beta = get_node_attr_f(node, "beta", 1)
                transA = get_node_attr_i(node, "transA", 0)
                transB = get_node_attr_i(node, "transB", 0)

                if alpha == 1 and beta == 1 and transA == 0 and transB == 1:
                    # InnerProduct-like A * B + C
                    layer.op_type = "InnerProduct"
                else:
                    layer.op_type = "Gemm"
            elif op in [
                "GlobalAveragePool",
                "GlobalMaxPool",
                "adaptive_avg_pool2d",
                "adaptive_max_pool2d",
            ]:
                layer.op_type = "Pooling"
            elif op == "GroupNorm":
                layer.op_type = "GroupNorm"
            elif op == "GRU":
                layer.op_type = "GRU"
            elif op == "HardSigmoid":
                layer.op_type = "HardSigmoid"
            elif op == "HardSwish":
                layer.op_type = "HardSwish"
            elif op == "ImageScaler":
                layer.op_type = "Scale"
            elif op == "InstanceNormalization":
                layer.op_type = "InstanceNorm"
            elif op == "LayerNorm":
                layer.op_type = "LayerNorm"
            elif op == "LeakyRelu" or op == "Relu":
                layer.op_type = "ReLU"
            elif op == "LRN":
                layer.op_type = "LRN"
            elif op == "LSTM":
                layer.op_type = "LSTM"
            elif op == "MatMul":
                if (
                    node.input[1] in self.weights
                    and len(self.weights[node.input[1]].dims) == 2
                ):
                    layer.op_type = "InnerProduct"
                else:
                    layer.op_type = "Gemm"
            elif op == "MultiHeadAttention":
                layer.op_type = "MultiHeadAttention"
            elif op == "Normalize":
                layer.op_type = "Normalize"
            elif op == "Pad":
                layer.op_type = "Padding"
            elif op == "PRelu":
                layer.op_type = "PReLU"
            elif op in [
                "ReduceMax",
                "ReduceMin",
                "ReduceMean",
                "ReduceProd",
                "ReduceSum",
                "ReduceSumSquare",
                "ReduceL1",
                "ReduceL2",
                "ReduceLogSum",
                "ReduceLogSumExp",
            ]:
                layer.op_type = "Reduction"
            elif op == "Reorg":
                layer.op_type = "Reorg"
            elif op == "Reshape":
                layer.op_type = "Reshape"
            elif op == "RNN":
                layer.op_type = "RNN"
            elif op == "ShuffleChannel":
                layer.op_type = "ShuffleChannel"
            elif op == "Sigmoid":
                layer.op_type = "Sigmoid"
            elif op == "SkipLayerNormalization":
                layer.op_type = "SkipLayerNormalization"
            elif op == "Softmax":
                layer.op_type = "Softmax"
            elif op == "Softplus":
                layer.op_type = "Softplus"
            elif op == "Split":
                layer.op_type = "Slice"
            elif op == "Squeeze":
                layer.op_type = "Squeeze"
            elif op == "Sum":
                layer.op_type = "Eltwise"
            elif op == "Swish":
                layer.op_type = "Swish"
            elif op == "Transpose":
                layer.op_type = "Permute"
            elif op == "Upsample" or op == "Resize":
                layer.op_type = "Interp"
            elif op == "Unsqueeze":
                layer.op_type = "ExpandDims"
            else:
                error_msg = f"{op} not currently supported by NCNN."
                raise ValueError(error_msg)

            layer.name = name
            layer.num_inputs = input_size
            layer.num_outputs = output_size
            layer.params.set_op(layer.op_type)

            for input_name in node.input:
                # check weight
                if input_name and not (
                    input_name in self.weights and self.node_reference[input_name] == 0
                ):
                    if input_name in split_node_reference:
                        refidx = split_node_reference[input_name] - 1
                        split_node_reference[input_name] = refidx
                        input_name = f"{input_name}_splitncnn_{refidx}"

                    layer.inputs.append(input_name)

            for o in range(output_size):
                layer.outputs.append(node.output[o])

            if op == "Abs":
                layer.add_param(0, UOT.ABS)
            elif op == "Acos":
                layer.add_param(0, UOT.ACOS)
            elif layer.op_type == "BinaryOp":
                if op == "Add":
                    layer.add_param(0, BOT.ADD)
                elif op == "Div":
                    layer.add_param(0, BOT.DIV)
                elif op == "Max":
                    layer.add_param(0, BOT.MAX)
                elif op == "Min":
                    layer.add_param(0, BOT.MIN)
                elif op == "Mul":
                    layer.add_param(0, BOT.MUL)
                elif op == "Pow":
                    layer.add_param(0, BOT.POW)
                elif op == "RDiv":
                    layer.add_param(0, BOT.RDIV)
                elif op == "RSub":
                    layer.add_param(0, BOT.RSUB)
                elif op == "Sub":
                    layer.add_param(0, BOT.SUB)

                with_scalar = get_node_attr_i(node, "with_scalar", 0)
                b = get_node_attr_f(node, "b", 0)
                if with_scalar:
                    layer.add_param(1, with_scalar)
                    layer.add_param(2, b)
            elif op == "Asin":
                layer.add_param(0, UOT.ASIN)
            elif op == "Atan":
                layer.add_param(0, UOT.ATAN)
            elif op == "AveragePool" or op == "MaxPool":
                auto_pad = get_node_attr_s(node, "auto_pad")
                ceil_mode = get_node_attr_i(node, "ceil_mode", 0)
                kernel_shape = get_node_attr_ai(node, "kernel_shape")
                strides = get_node_attr_ai(node, "strides")
                pads = get_node_attr_ai(node, "pads")

                pool = int(op == "AveragePool")

                if ceil_mode == 1:
                    pad_mode = PAM.FULL
                elif auto_pad == "SAME_UPPER":
                    pad_mode = PAM.SAMEUPPER
                elif auto_pad == "SAME_LOWER":
                    pad_mode = PAM.SAMELOWER
                else:
                    pad_mode = PAM.VALID

                layer.add_param(0, pool)

                if kernel_shape.size == 1:
                    layer.add_param(1, int(kernel_shape[0]))
                elif kernel_shape.size == 2:
                    layer.add_param(1, int(kernel_shape[1]))
                    layer.add_param(11, int(kernel_shape[0]))

                if strides.size == 1:
                    layer.add_param(2, int(strides[0]))
                elif strides.size == 2:
                    layer.add_param(2, int(strides[1]))
                    layer.add_param(12, int(strides[0]))

                if pads.size == 1:
                    layer.add_param(3, int(pads[0]))
                elif pads.size == 2:
                    layer.add_param(3, int(pads[1]))
                    layer.add_param(13, int(pads[0]))
                elif pads.size == 4:
                    layer.add_param(3, int(pads[1]))
                    layer.add_param(13, int(pads[0]))
                    layer.add_param(14, int(pads[3]))
                    layer.add_param(15, int(pads[2]))

                layer.add_param(5, pad_mode)

                if pool:
                    avgpool_count_include_pad = get_node_attr_i(
                        node, "count_include_pad", 0
                    )
                    layer.add_param(6, avgpool_count_include_pad)
            elif op == "BatchNormalization":
                epsilon = get_node_attr_f(node, "epsilon", 0.00001)
                scale = self.weights[node.input[1]]
                B = self.weights[node.input[2]]
                mean = self.weights[node.input[3]]
                var = self.weights[node.input[4]]
                channels = get_tensor_proto_data_size(scale, scale.data_type)

                layer.add_param(0, channels)

                bin_length += self.add_weight(layer, "slope", scale)
                bin_length += self.add_weight(layer, "mean", mean)

                # apply epsilon to var
                v = onph.to_array(var)
                ve = np.array([v[i] + epsilon for i in range(channels)], np.float32)
                bin_length += self.add_weight(layer, "variance", ve)
                bin_length += self.add_weight(layer, "bias", B)
            elif op == "BiasGelu":
                B = self.weights[node.input[1]]

                layer.add_param(0, get_tensor_proto_data_size(B, B.data_type))

                bin_length += self.add_weight(layer, "bias", B)
            elif op == "Ceil":
                layer.add_param(0, UOT.CEIL)
            elif op == "Clip":
                if len(node.input) == 1:
                    minimum = get_node_attr_f(node, "min", -FLOAT32_MAX)
                    maximum = get_node_attr_f(node, "max", FLOAT32_MAX)
                else:
                    minimum = (
                        get_node_attr_from_input_f(self.weights[node.input[1]])
                        if node.input[1] in self.weights
                        else -FLOAT32_MAX
                    )
                    maximum = (
                        get_node_attr_from_input_f(self.weights[node.input[2]])
                        if node.input[2] in self.weights
                        else FLOAT32_MAX
                    )

                layer.add_param(0, minimum)
                layer.add_param(1, maximum)
            elif op == "Concat":
                axis = get_node_attr_i(node, "axis", 1)
                layer.add_param(0, axis - 1 if axis > 0 else axis)
            elif op == "Constant":
                logger.error("chaiNNer: code should not have reached inside Constant")
            elif op == "Conv":
                W = self.weights[node.input[1]]

                num_filter = W.dims[0]
                has_bias = int(len(node.input) == 3)

                auto_pad = get_node_attr_s(node, "auto_pad")
                kernel_shape = get_node_attr_ai(node, "kernel_shape")
                dilations = get_node_attr_ai(node, "dilations")
                strides = get_node_attr_ai(node, "strides")
                pads = get_node_attr_ai(node, "pads")
                group = get_node_attr_i(node, "group", 1)

                layer.add_param(0, num_filter)

                if kernel_shape.size == 1:
                    layer.add_param(1, int(kernel_shape[0]))
                elif kernel_shape.size == 2:
                    layer.add_param(1, int(kernel_shape[1]))
                    layer.add_param(11, int(kernel_shape[0]))

                if dilations.size == 1:
                    layer.add_param(2, int(dilations[0]))
                elif dilations.size == 2:
                    layer.add_param(2, int(dilations[1]))
                    layer.add_param(12, int(dilations[0]))

                if strides.size == 1:
                    layer.add_param(3, int(strides[0]))
                elif strides.size == 2:
                    layer.add_param(3, int(strides[1]))
                    layer.add_param(13, int(strides[0]))

                if auto_pad == "SAME_UPPER":
                    layer.add_param(4, -233)
                elif auto_pad == "SAME_LOWER":
                    layer.add_param(4, -234)
                else:
                    if pads.size == 1:
                        layer.add_param(4, int(pads[0]))
                    elif pads.size == 2:
                        layer.add_param(4, int(pads[1]))
                        layer.add_param(14, int(pads[0]))
                    elif pads.size == 4:
                        layer.add_param(4, int(pads[1]))
                        layer.add_param(14, int(pads[0]))
                        layer.add_param(15, int(pads[3]))
                        layer.add_param(16, int(pads[2]))

                layer.add_param(5, has_bias)

                layer.add_param(6, get_tensor_proto_data_size(W, W.data_type))

                if group > 1:
                    layer.add_param(7, int(group))

                quantize_tag = DTYPE_FP16 if is_fp16 else DTYPE_FP32
                bin_length += self.add_weight(layer, "weight", W, quantize_tag)

                if has_bias:
                    B = self.weights[node.input[2]]
                    bin_length += self.add_weight(layer, "bias", B)
            elif op == "ConvTranspose":
                W = self.weights[node.input[1]]

                has_bias = int(len(node.input) == 3)

                auto_pad = get_node_attr_s(node, "auto_pad")
                kernel_shape = get_node_attr_ai(node, "kernel_shape")
                dilations = get_node_attr_ai(node, "dilations")
                strides = get_node_attr_ai(node, "strides")
                output_padding = get_node_attr_ai(node, "output_padding")
                output_shape = get_node_attr_ai(node, "output_shape")
                pads = get_node_attr_ai(node, "pads")
                group = get_node_attr_i(node, "group", 1)
                num_filter = W.dims[1] * group

                layer.add_param(0, num_filter)

                if kernel_shape.size == 1:
                    layer.add_param(1, int(kernel_shape[0]))
                elif kernel_shape.size == 2:
                    layer.add_param(1, int(kernel_shape[1]))
                    layer.add_param(11, int(kernel_shape[0]))

                if dilations.size == 1:
                    layer.add_param(2, int(dilations[0]))
                elif dilations.size == 2:
                    layer.add_param(2, int(dilations[1]))
                    layer.add_param(12, int(dilations[0]))

                if strides.size == 1:
                    layer.add_param(3, int(strides[0]))
                elif strides.size == 2:
                    layer.add_param(3, int(strides[1]))
                    layer.add_param(13, int(strides[0]))

                if auto_pad == "SAME_UPPER":
                    layer.add_param(4, -233)
                elif auto_pad == "SAME_LOWER":
                    layer.add_param(4, -234)
                else:
                    if pads.size == 1:
                        layer.add_param(4, int(pads[0]))
                    elif pads.size == 2:
                        layer.add_param(4, int(pads[1]))
                        layer.add_param(14, int(pads[0]))
                    elif pads.size == 4:
                        layer.add_param(4, int(pads[1]))
                        layer.add_param(14, int(pads[0]))
                        layer.add_param(15, int(pads[3]))
                        layer.add_param(16, int(pads[2]))

                if output_padding.size == 1:
                    layer.add_param(18, int(output_padding[0]))
                elif output_padding.size == 2:
                    layer.add_param(18, int(output_padding[1]))
                    layer.add_param(19, int(output_padding[0]))

                if output_shape.size == 1:
                    layer.add_param(20, int(output_shape[0]))
                elif output_shape == 2:
                    layer.add_param(20, int(output_shape[1]))
                    layer.add_param(21, int(output_shape[0]))

                layer.add_param(5, has_bias)

                weight_data_size = get_tensor_proto_data_size(W, W.data_type)
                layer.add_param(6, weight_data_size)

                if group > 1:
                    layer.add_param(7, group)

                quantize_tag = DTYPE_FP16 if is_fp16 else DTYPE_FP32
                weight_data = onph.to_array(W)
                bin_length += self.add_weight(
                    layer, "weight", weight_data.swapaxes(0, 1), quantize_tag
                )

                if has_bias:
                    B = self.weights[node.input[2]]
                    bin_length += self.add_weight(layer, "bias", B)
            elif op == "Cos":
                layer.add_param(0, UOT.COS)
            elif op == "Crop":
                starts = get_node_attr_ai(node, "starts")
                layer.add_param(9, [starts.size, *starts])

                ends = get_node_attr_ai(node, "ends")
                layer.add_param(10, [ends.size, *ends])

                axes = get_node_attr_ai(node, "axis")
                layer.add_param(11, [axes.size, *axes])
            elif op == "DepthToSpace":
                # pixelshuffle
                scale_factor = get_node_attr_i(node, "blocksize", 1)
                mode = get_node_attr_s(node, "mode")
                layer.add_param(0, scale_factor)
                if mode == "CRD":
                    layer.add_param(1, 0)
                elif mode == "DCR":
                    layer.add_param(1, 1)
            elif op == "Dropout":
                pass
            elif op == "Elu":
                alpha = get_node_attr_f(node, "alpha", 1)
                layer.add_param(0, alpha)
            elif op == "EmbedLayerNormalization":
                logger.error(f"chaiNNer: No NCNN documentation for {op} yet, will not function")
                words = self.weights[node.input[2]]
                positions = self.weights[node.input[3]]
                W = self.weights[node.input[5]]
                B = self.weights[node.input[6]]

                layer.add_param(0, get_tensor_proto_data_size(B, B.data_type))
                layer.add_param(1, get_tensor_proto_data_size(words, words.data_type))
                layer.add_param(
                    2, get_tensor_proto_data_size(positions, positions.data_type)
                )

                quantize_tag = DTYPE_FP16 if is_fp16 else DTYPE_FP32
                bin_length += self.add_weight(layer, "words", words, DTYPE_FP32)
                bin_length += self.add_weight(layer, "positions", positions, DTYPE_FP32)
                bin_length += self.add_weight(layer, "weight", W, quantize_tag)
                bin_length += self.add_weight(layer, "bias", B)
            elif op == "Exp":
                layer.add_param(0, UOT.EXP)
            elif op == "Flatten":
                axis = get_node_attr_i(node, "axis", 1)
                if axis != 1:
                    raise ValueError(f"Unsupported Flatten axis {axis}.")
            elif op == "Floor":
                layer.add_param(0, UOT.FLOOR)
            elif op == "Gelu":
                layer.add_param(0, 1)
            elif op == "Gemm":
                alpha = get_node_attr_f(node, "alpha", 1)
                beta = get_node_attr_f(node, "beta", 1)
                transA = get_node_attr_i(node, "transA", 0)
                transB = get_node_attr_i(node, "transB", 0)

                if alpha == 1 and beta == 1 and transA == 0 and transB == 1:
                    # InnerProduct-like A * B * C
                    B = self.weights[node.input[1]]
                    C = self.weights[node.input[2]]

                    layer.add_param(0, get_tensor_proto_data_size(C, C.data_type))
                    layer.add_param(1, 1)
                    layer.add_param(2, get_tensor_proto_data_size(B, B.data_type))

                    quantize_tag = DTYPE_FP16 if is_fp16 else DTYPE_FP32
                    bin_length += self.add_weight(layer, "B", B, quantize_tag)
                    bin_length += self.add_weight(layer, "C", C)
                else:
                    # gemm
                    layer.add_param(0, alpha)
                    layer.add_param(1, beta)
                    layer.add_param(2, transA)
                    layer.add_param(3, transB)
            elif op == "GlobalAveragePool" or op == "GlobalMaxPool":
                layer.add_param(0, int(op == "GlobalAveragePool"))
                layer.add_param(4, 1)
            elif op == "adaptive_avg_pool2d" or op == "adaptive_max_pool2d":
                out_shape_tp = self.weights[node.input[1]]
                out_shape = get_node_attr_from_input_ai(out_shape_tp)

                layer.add_param(0, int(op == "adaptive_avg_pool2d"))
                layer.add_param(7, 1)
                if out_shape.size == 1:
                    layer.add_param(8, int(out_shape[0]))
                elif out_shape.size == 2:
                    layer.add_param(8, int(out_shape[1]))  # out_w
                    layer.add_param(18, int(out_shape[0]))  # out_h
            elif op == "GroupNorm":
                groups = get_node_attr_i(node, "groups", 1)
                channels = get_node_attr_i(node, "channels", 1)
                eps = get_node_attr_f(node, "epsilon", 0.00001)
                affine = get_node_attr_i(node, "affine", 1)

                if affine:
                    # discard affine-less S=1 B=0
                    affine_S = get_node_attr_from_input_af(self.weights[node.input[1]])
                    affine_B = get_node_attr_from_input_af(self.weights[node.input[2]])
                    if (
                        affine_S.size == 1
                        and affine_S[0] == 1
                        and affine_B.size == 1
                        and affine_B[0] == 0
                    ):
                        affine = 0
                    else:
                        if np.any(affine_S[:channels] != 1) or np.any(
                            affine_B[:channels] != 0
                        ):
                            affine = 1
                        else:
                            affine = 0

                layer.add_param(0, groups)
                layer.add_param(1, channels)
                layer.add_param(2, eps)
                layer.add_param(3, affine)
                if affine:
                    scale = self.weights[node.input[1]]
                    B = self.weights[node.input[2]]

                    bin_length += self.add_weight(layer, "scale", scale)
                    bin_length += self.add_weight(layer, "bias", B)
            elif op == "GRU":
                # W = self.weights[node.input[1]]
                # R = self.weights[node.input[2]]
                # B = self.weights[node.input[3]]

                # hidden_size = get_node_attr_i(node, "hidden_size", 0)
                # direction = get_node_attr_s(node, "direction")

                # if direction == "forward":
                #    direction_type = GRU.FORWARD
                # elif direction == "reverse":
                #    direction_type = GRU.REVERSE
                # elif direction == "bidirectional":
                #    direction_type = GRU.BIDIRECTIONAL

                # weight_data_size = get_tensor_proto_data_size(W)

                # layer.add_param(0, hidden_size)
                # layer.add_param(1, weight_data_size)
                # layer.add_param(2, direction_type)

                # num_directions = 2 if direction_type == GRU.BIDIRECTIONAL else 1

                # reorder num_directions-URN-hidden_size to num_directions-RUN-hidden_size
                # quantize_tag = DTYPE_FP16 if is_fp16 else DTYPE_FP32

                # logger.error(
                #    "Not sure GRU weight reordering is accurate, "
                #    "docs and code comments appear to give different shape orders"
                # )

                # W_array = onph.to_array(W)
                # W_array = np.stack(
                #    (W_array[:, 1, :], W_array[:, 0, :], W_array[:, 2, :]), axis=1
                # )
                # bin_length += self.add_weight(layer, W_array, "weight_xc_data", quantize_tag, is_fp16)

                # reduce U and R bias except N
                # reorder num_directions-URN-hidden to num_directions-RUN-hidden
                # B_array = onph.to_array(B)

                # bias_data_size_g = B_array.size / 6 / num_directions
                # for i in range(bias_data_size_g)[1:]:
                #    pass
                raise RuntimeError(
                    "GRU not implemented yet, please report issue with model used"
                )
            elif op == "HardSigmoid" or op == "Hard Swish":
                alpha = get_node_attr_f(node, "alpha", 0.2)
                beta = get_node_attr_f(node, "beta", 0.5)

                layer.add_param(0, alpha)
                layer.add_param(1, beta)
            elif op == "ImageScaler":
                bias = get_node_attr_af(node, "bias")
                scale = get_node_attr_f(node, "scale", 1)
                channels = bias.size

                layer.add_param(0, channels)
                layer.add_param(1, 1)

                bin_length += self.add_weight(layer, "scale", np.array((scale,) * 3))
                bin_length += self.add_weight(layer, "bias", bias)
            elif op == "InstanceNormalization":
                eps = get_node_attr_f(node, "epsilon", 0.00001)

                # Discard affine-less S=1 B=0
                affine_S = get_node_attr_from_input_af(self.weights[node.input[1]])
                affine_B = get_node_attr_from_input_af(self.weights[node.input[2]])
                channels = affine_S.size

                if np.any(affine_S[:channels] != 1) or np.any(affine_B[:channels] != 0):
                    affine = 1
                else:
                    affine = 0

                layer.add_param(0, channels)
                layer.add_param(1, eps)
                layer.add_param(2, affine)
                if affine:
                    scale = self.weights[node.input[1]]
                    B = self.weights[node.input[2]]

                    bin_length += self.add_weight(layer, "scale", scale)
                    bin_length += self.add_weight(layer, "bias", B)
            elif op == "LayerNorm":
                eps = get_node_attr_f(node, "epsilon", 0.00001)
                affine = get_node_attr_i(node, "affine", 1)

                if affine:
                    # discard affine-less S=1 B=0
                    affine_S = get_node_attr_from_input_af(self.weights[node.input[1]])
                    affine_B = get_node_attr_from_input_af(self.weights[node.input[2]])
                    affine_size = affine_S.size

                    if np.any(affine_S[:affine_size] != 1) or np.any(
                        affine_B[:affine_size]
                    ):
                        affine = 1
                    else:
                        affine = 0

                    if affine:
                        layer.add_param(0, affine_size)

                layer.add_param(1, eps)
                layer.add_param(2, affine)

                if affine:
                    scale = self.weights[node.input[1]]
                    B = self.weights[node.input[2]]

                    bin_length += self.add_weight(layer, "scale", scale)
                    bin_length += self.add_weight(layer, "bias", B)
            elif op == "LeakyRelu":
                alpha = get_node_attr_f(node, "alpha", 0.01)
                layer.add_param(0, alpha)
            elif op == "Log":
                layer.add_param(0, UOT.LOG)
            elif op == "LRN":
                layer.add_param(0, 0)
                layer.add_param(1, get_node_attr_i(node, "size", 1))
                layer.add_param(2, get_node_attr_f(node, "alpha", 1))
                layer.add_param(3, get_node_attr_f(node, "beta", 0.5))
                layer.add_param(4, get_node_attr_f(node, "bias", 1))
            elif op == "LSTM":
                # W = self.weights[node.input[1]]
                # R = self.weights[node.input[2]]
                # B = self.weights[node.input[3]]

                # hidden_size = get_node_attr_i(node, "hidden_size", 0)
                # direction = get_node_attr_s(node, "direction")

                # if direction == "forward":
                #    direction_type = GRU.FORWARD
                # elif direction == "reverse":
                #    direction_type = GRU.REVERSE
                # elif direction  == "bidirectional":
                #    direction_type = GRU.BIDIRECTIONAL
                raise RuntimeError(
                    "LSTM not implemented yet, please report issue with model used"
                )
            elif op == "MatMul":
                if node.input[1] in self.weights:
                    # InnerProduct
                    B = self.weights[node.input[1]]
                    weight_data_size = get_tensor_proto_data_size(B, B.data_type)
                    num_output = B.dims[-1]

                    layer.add_param(0, num_output)
                    layer.add_param(1, 0)
                    layer.add_param(2, weight_data_size)

                    B_array = onph.to_array(B)
                    bin_length += self.add_weight(layer, "bias", B_array.T, DTYPE_FP32)
                # There is a dead else here, not sure if this was incomplete code
            elif op == "MultiHeadAttention":
                # embed_dim = get_node_attr_i(node, "embed_dim", 0)
                # num_heads = get_node_attr_i(node, "num_heads", 0)

                # layer.add_param(0, embed_dim)
                # layer.add_param(1, num_heads)

                # if len(node.input) == 5:
                #    qkvw = self.weights[node.input[1]]
                #    qkvb = self.weights[node.input[2]]
                #    ow = self.weights[node.input[3]]
                #    ob = self.weights[node.input[4]]

                #    weight_data_size = get_tensor_proto_data_size(ow)

                #    layer.add_param(2, weight_data_size)

                #    quantize_tag = DTYPE_FP16 if is_fp16 else DTYPE_FP32
                raise RuntimeError(
                    "MultiHeadAttention not implemented, please report issue with model used"
                )
            elif op == "Neg":
                layer.add_param(0, UOT.NEG)
            elif op == "Normalize":
                eps = get_node_attr_f(node, "eps", 0)

                layer.add_param(1, 1)  # channel_shared
                layer.add_param(2, eps)
                layer.add_param(3, 1)  # scale_data_size
                layer.add_param(9, NEM.PYTORCH)

                bin_length += self.add_weight(layer, "scale", 1)
            elif op == "Pad":
                mode = get_node_attr_s(node, "mode")
                value = get_node_attr_f(node, "value", 0)

                if len(node.input) == 1:
                    pads = get_node_attr_ai(node, "pads")
                else:
                    pads = get_node_attr_from_input_ai(self.weights[node.input[1]])

                if mode == "edge":
                    ptype = PAT.REPLICATE
                elif mode == "reflect":
                    ptype = PAT.REFLECT
                else:
                    ptype = PAT.CONSTANT

                pad_size = pads.size
                top = bottom = front = behind = 0
                if pad_size == 8:
                    # NCHW
                    top = pads[2]
                    bottom = pads[6]
                    left = pads[3]
                    right = pads[7]
                    front = pads[1]
                    behind = pads[5]
                elif pad_size == 6:
                    # NHW
                    top = pads[1]
                    bottom = pads[4]
                    left = pads[2]
                    right = pads[5]
                else:
                    # NW
                    left = pads[1]
                    right = pads[3]

                layer.add_param(0, int(top))
                layer.add_param(1, int(bottom))
                layer.add_param(2, int(left))
                layer.add_param(3, int(right))
                layer.add_param(4, int(ptype))
                layer.add_param(5, int(value))
                layer.add_param(7, int(front))
                layer.add_param(8, int(behind))
            elif op == "PixelShuffle":
                layer.add_param(0, get_node_attr_i(node, "scale_factor", 1))
            elif op == "PRelu":
                slope = self.weights[node.input[1]]
                num_slope = get_tensor_proto_data_size(slope, slope.data_type)

                layer.add_param(0, num_slope)

                bin_length += self.add_weight(layer, "slope", slope)
            elif op == "Reciprocal":
                layer.add_param(0, UOT.RECIPROCAL)
            elif op in [
                "ReduceMax",
                "ReduceMin",
                "ReduceMean",
                "ReduceProd",
                "ReduceSum",
                "ReduceSumSquare",
                "ReduceL1",
                "ReduceL2",
                "ReduceLogSum",
                "ReduceLogSumExp",
            ]:
                if op == "ReduceSum":
                    op_type = ROT.SUM
                elif op == "ReduceSumSquare":
                    op_type = ROT.SUMSQ
                elif op == "ReduceMean":
                    op_type = ROT.MEAN
                elif op == "ReduceMax":
                    op_type = ROT.MAX
                elif op == "ReduceMin":
                    op_type = ROT.MIN
                elif op == "ReduceProd":
                    op_type = ROT.PROD
                elif op == "ReduceL1":
                    op_type = ROT.L1
                elif op == "ReduceL2":
                    op_type = ROT.L2
                elif op == "ReduceLogSum":
                    op_type = ROT.LOGSUM
                elif op == "ReduceLogSumExp":
                    op_type = ROT.LOGSUMEXP
                else:
                    op_type = -233

                layer.add_param(0, op_type)

                axes = get_node_attr_ai(node, "axes")
                keepdims = get_node_attr_i(node, "keepdims", 1)

                if axes.size > 0:
                    # if axes set, reduce according to axes
                    layer.add_param(1, 0)

                    for axis in axes:
                        if axis == 0 or axis > 4 or axis < -3:
                            raise ValueError(f"Unsupported axis {axis} in Reduction")
                    layer.add_param(
                        3,
                        [axes.size, *[a - 1 if a > 0 else a for a in axes]],
                    )
                else:
                    # if axes not set, reduce all axes by default
                    layer.add_param(1, 1)

                layer.add_param(4, keepdims)
                logger.error("chaiNNer: No NCNN documentation for Reduction param 5")
                layer.add_param(5, 1)
            elif op == "Reorg":
                layer.add_param(0, get_node_attr_i(node, "stride", 1))
            elif op == "Reshape":
                if len(node.input) == 1:
                    shape = get_node_attr_ai(node, "shape")
                else:
                    shape = get_node_attr_from_input_ai(self.weights[node.input[1]])

                shape_size = shape.size
                if shape_size == 1:
                    logger.error("chaiNNer: Should never reach shape.size == 1 in Reshape")
                    layer.add_param(0, int(shape[0]))
                elif shape_size == 2:
                    layer.add_param(0, int(shape[1]))
                elif shape_size == 3:
                    layer.add_param(0, int(shape[2]))
                    layer.add_param(1, int(shape[1]))
                elif shape_size == 4:
                    layer.add_param(0, int(shape[3]))
                    layer.add_param(1, int(shape[2]))
                    layer.add_param(2, int(shape[1]))
                elif shape_size == 5:
                    layer.add_param(0, int(shape[3] * shape[3]))
                    layer.add_param(1, int(shape[2]))
                    layer.add_param(2, int(shape[1]))
            elif op == "Resize":
                mode = get_node_attr_s(node, "mode")
                align = get_node_attr_s(node, "coordinate_transformation_mode")

                if len(node.input) == 2:
                    # opset 10
                    scales = get_node_attr_from_input_af(self.weights[node.input[1]])
                    sizes = np.empty(0, np.int32)
                else:
                    # opset 11+
                    scales = get_node_attr_from_input_af(self.weights[node.input[2]])
                    if len(node.input) >= 4:
                        sizes = get_node_attr_from_input_ai(self.weights[node.input[3]])
                    else:
                        sizes = np.empty(0, np.int32)

                if mode == "linear":
                    resize_type = IRT.BILINEAR
                elif mode == "cubic":
                    resize_type = IRT.BICUBIC
                else:
                    resize_type = IRT.NEAREST

                if scales.size == 0 and sizes.size == 0:
                    raise ValueError(
                        "Unsupported Resize scales and sizes are all empty."
                    )

                if scales.size == 2:
                    h_scale = 1
                    w_scale = scales[1]
                elif scales.size == 3:
                    h_scale = scales[1]
                    w_scale = scales[2]
                elif scales.size == 4:
                    if scales[1] != 1:
                        raise TypeError(f"Unsupported Resize scales {scales}.")
                    h_scale = scales[2]
                    w_scale = scales[3]
                else:
                    h_scale = 1
                    w_scale = 1

                if sizes.size == 2:
                    output_height = 0
                    output_width = sizes[1]
                elif sizes.size == 3:
                    output_height = sizes[1]
                    output_width = sizes[2]
                elif sizes.size == 4:
                    output_height = sizes[2]
                    output_width = sizes[3]
                else:
                    output_height = 0
                    output_width = 0

                align_corner = int(align == "align_corners")

                layer.add_param(0, resize_type)
                layer.add_param(1, float(h_scale))
                layer.add_param(2, float(w_scale))
                layer.add_param(3, int(output_height))
                layer.add_param(4, int(output_width))
                layer.add_param(6, align_corner)
            elif op == "RNN":
                W = self.weights[node.input[1]]
                R = self.weights[node.input[2]]
                B = self.weights[node.input[3]]

                hidden_size = get_node_attr_i(node, "hidden_size", 0)
                direction = get_node_attr_s(node, "direction")

                if direction == "reverse":
                    direction_type = GRU.REVERSE
                elif direction == "bidirectional":
                    direction_type = GRU.BIDIRECTIONAL
                else:
                    direction_type = GRU.FORWARD

                weight_data_size = get_tensor_proto_data_size(W, W.data_type)

                layer.add_param(0, hidden_size)
                layer.add_param(1, weight_data_size)
                layer.add_param(2, direction_type)

                quantize_tag = DTYPE_FP16 if is_fp16 else DTYPE_FP32
                bin_length += self.add_weight(layer, "weight", W, quantize_tag)

                # reduce xc and hc bias
                reduced_B = np.sum(onph.to_array(B), 1)
                bin_length += self.add_weight(layer, "bias", reduced_B, quantize_tag)

                bin_length += self.add_weight(layer, "R", R, quantize_tag)
            elif op == "ShuffleChannel":
                layer.add_param(0, get_node_attr_i(node, "group", 1))
                layer.add_param(1, get_node_attr_i(node, "reverse", 0))
            elif op == "Sigmoid":
                pass
            elif op == "Sin":
                layer.add_param(0, UOT.SIN)
            elif op == "SkipLayerNormalization":
                logger.error(f"chaiNNer: No NCNN documentation for {op} yet, will not function")
                W = self.weights[node.input[2]]
                B = self.weights[node.input[3]]
                B2 = self.weights[node.input[4]]

                layer.add_param(0, get_tensor_proto_data_size(B, B.data_type))

                quantize_tag = DTYPE_FP16 if is_fp16 else DTYPE_FP32
                bin_length += self.add_weight(layer, "weight", W, quantize_tag)
                bin_length += self.add_weight(layer, "bias1", B, DTYPE_FP32)
                bin_length += self.add_weight(layer, "bias2", B2, DTYPE_FP32)
            elif op == "Slice":
                input_size = len(node.input)
                if input_size == 1:
                    starts = get_node_attr_ai(node, "starts")
                    ends = get_node_attr_ai(node, "ends")
                    axes = get_node_attr_ai(node, "axes")
                    steps = get_node_attr_ai(node, "steps")
                else:
                    starts = get_node_attr_from_input_ai(self.weights[node.input[1]])
                    ends = get_node_attr_from_input_ai(self.weights[node.input[2]])
                    if input_size >= 4:
                        axes = get_node_attr_from_input_ai(self.weights[node.input[3]])
                    else:
                        axes = np.empty(0, np.int32)
                    if input_size >= 5:
                        steps = get_node_attr_from_input_ai(self.weights[node.input[4]])
                    else:
                        steps = np.empty(0, np.int32)

                assert np.all(steps != 1), f"Unsupported Slice step {steps}"

                # Filter out N-dim axis
                if axes.size:
                    for i, axis in enumerate(axes):
                        if axis == 0:
                            np.delete(starts, i)
                            np.delete(ends, i)
                            np.delete(axes, i)
                            break

                layer.add_param(9, [starts.size, *list(starts)])
                layer.add_param(10, [ends.size, *list(ends)])
                if axes.size:
                    assert np.all(
                        axes != 0 and axes <= 3 and axes >= -3
                    ), f"Unsupported Slice axes {axes}"
                    layer.add_param(
                        11, [axes.size, *[a - 1 if a > 0 else a for a in axes]]
                    )
            elif op == "Softmax":
                axis = get_node_attr_i(node, "axis", 1)
                layer.add_param(0, axis - 1)
                layer.add_param(1, 1)
            elif op == "Split":
                axis = get_node_attr_i(node, "axis", 0)
                splits = get_node_attr_ai(node, "split")

                assert axis >= 1, f"Unsupported axis {axis} in Split"

                if splits.size:
                    layer.add_param(0, [output_size, *list(splits[:-1]), -233])
                else:
                    layer.add_param(
                        0, [output_size, *[-233 for _ in range(output_size)]]
                    )
                layer.add_param(1, axis - 1)
            elif op == "Sqrt":
                layer.add_param(0, UOT.SQRT)
            elif op == "Squeeze":
                axes = get_node_attr_ai(node, "axes")

                if axes.size:
                    assert np.all(
                        axes != 0 and axes <= 4 and axes >= -3
                    ), f"Unsupported Squeeze axes {axes}"

                    layer.add_param(
                        3, [axes.size, *[a - 1 if a > 0 else a for a in axes]]
                    )
                else:
                    layer.add_param(0, 1)
                    layer.add_param(1, 1)
                    layer.add_param(2, 1)
            elif op == "Sum":
                layer.add_param(0, EOT.SUM)
            elif op == "Swish":
                pass
            elif op == "Tan":
                layer.add_param(0, UOT.TAN)
            elif op == "Tanh":
                layer.add_param(0, UOT.TANH)
            elif op == "Transpose":
                perm = get_node_attr_ai(node, "perm")
                if perm.size == 3:
                    if (perm[1] == 1 and perm[2] == 2) or (
                        perm[0] == 1 and perm[1] == 0 and perm[2] == 2
                    ):
                        layer.add_param(0, POT.WH_WHC_WHDC)
                    elif (perm[1] == 2 and perm[2] == 1) or (
                        perm[0] == 2 and perm[1] == 0 and perm[2] == 1
                    ):
                        layer.add_param(0, POT.HW_HWC_HWDC)
                elif perm.size == 4:
                    if perm[1] == 1 and perm[2] == 2 and perm[3] == 3:
                        layer.add_param(0, POT.WH_WHC_WHDC)
                    elif perm[1] == 1 and perm[2] == 3 and perm[3] == 2:
                        layer.add_param(0, POT.HW_HWC_HWDC)
                    elif perm[1] == 2 and perm[2] == 1 and perm[3] == 3:
                        layer.add_param(0, POT.WCH_WDHC)
                    elif perm[1] == 2 and perm[2] == 3 and perm[3] == 1:
                        layer.add_param(0, POT.CWH_DWHC)
                    elif perm[1] == 3 and perm[2] == 1 and perm[3] == 2:
                        layer.add_param(0, POT.HCW_HDWC)
                    elif perm[1] == 3 and perm[2] == 2 and perm[3] == 1:
                        layer.add_param(0, POT.CHW_DHWC)
                elif perm.size == 5:
                    if perm[1] == 1 and perm[2] == 2 and perm[3] == 3 and perm[4] == 4:
                        layer.add_param(0, POT.WH_WHC_WHDC)
                    elif (
                        perm[1] == 1 and perm[2] == 3 and perm[3] == 4 and perm[4] == 2
                    ):
                        layer.add_param(0, POT.HW_HWC_HWDC)
                    elif (
                        perm[1] == 2 and perm[2] == 1 and perm[3] == 3 and perm[4] == 4
                    ):
                        layer.add_param(0, POT.WCH_WDHC)
                    elif (
                        perm[1] == 2 and perm[2] == 3 and perm[3] == 4 and perm[4] == 1
                    ):
                        layer.add_param(0, POT.CWH_DWHC)
                    elif (
                        perm[1] == 3 and perm[2] == 4 and perm[3] == 1 and perm[4] == 2
                    ):
                        layer.add_param(0, POT.HCW_HDWC)
                    elif (
                        perm[1] == 3 and perm[2] == 4 and perm[3] == 2 and perm[4] == 1
                    ):
                        layer.add_param(0, POT.CHW_DHWC)
                    else:
                        error_msg = f"Unsupported Transpose type {perm}"
                        raise ValueError(error_msg)
            elif op == "Upsample":
                mode = get_node_attr_s(node, "mode")
                align = get_node_attr_s(node, "coordinate_transformation_mode")

                if len(node.input) == 1:
                    scales = get_node_attr_af(node, "scales")
                else:
                    scales = get_node_attr_from_input_af(self.weights[node.input[1]])

                if mode == "bilinear" or mode == "linear":
                    resize_type = IRT.BILINEAR
                elif mode == "trilinear":
                    raise ValueError("Upsample does not support trilinear mode")
                else:
                    resize_type = IRT.NEAREST

                if scales.size == 2:
                    h_scale = 1
                    w_scale = scales[1]
                elif scales.size == 3:
                    h_scale = scales[1]
                    w_scale = scales[2]
                elif scales.size == 4:
                    h_scale = scales[2]
                    w_scale = scales[3]

                    if scales[1] != 1:
                        error_msg = f"Unsupported Upsample scales {scales}"
                        raise ValueError(error_msg)
                else:
                    error_msg = f"Unsupported Upsample scales {scales}"
                    raise ValueError(error_msg)

                align_corner = int(align == "align_corners")

                layer.add_param(0, resize_type)
                layer.add_param(1, float(h_scale))
                layer.add_param(2, float(w_scale))
                layer.add_param(6, align_corner)
            elif op == "Unsqueeze":
                axes = get_node_attr_ai(node, "axes")

                assert (
                    np.all(axes != 0) and np.all(axes <= 4) and np.all(axes >= -4)
                ), f"Unsupported axes {axes} in Unsqueeze"

                layer.add_param(
                    3, [axes.size, *[axis - 1 if axis > 0 else axis for axis in axes]]
                )
            else:
                # NCNN TODO: op specific param
                # This is presumably to catch anything they haven't written an op for yet
                for attr in node.attribute:
                    if attr.type == 1:
                        error_msg = f"Op {op} does not exist yet; {attr.name}={attr.f}"
                    elif attr.type == 2:
                        error_msg = f"Op {op} does not exist yet; {attr.name}={attr.i}"
                    elif attr.type == 3:
                        error_msg = f"Op {op} does not exist yet; {attr.name}={attr.s}"
                    else:
                        error_msg = (
                            f"Op {op} does not exist yet; {attr.name}={attr.type}"
                        )

                    raise ValueError(error_msg)

            ncnn_model.add_layer(layer)

            for o in range(output_size):
                output_name = node.output[o]
                if output_name in self.node_reference:
                    refcount = self.node_reference[output_name]
                    if refcount > 1:
                        ncnn_model.add_layer(
                            NcnnLayer(
                                "Split",
                                f"splitncnn_{internal_split}",
                                1,
                                refcount,
                                [output_name],
                                [
                                    f"{output_name}_splitncnn_{j}"
                                    for j in range(refcount)
                                ],
                            )
                        )

                        internal_split += 1

        ncnn_model.bin_length = bin_length
        NcnnOptimizer(ncnn_model).optimize()

        return ncnn_model

nodes/impl/onnx/onnx_to_ncnn.py
vladmandic/sd-extension-chainner
from __future__ import annotations
from weakref import WeakKeyDictionary
import onnxruntime as ort


def create_inference_session(
    model,
    gpu_index: int,
    execution_provider: str,
    should_tensorrt_fp16: bool = False,
    tensorrt_cache_path: str | None = None,
) -> ort.InferenceSession:
    if execution_provider == "TensorrtExecutionProvider":
        providers = [
            (
                "TensorrtExecutionProvider",
                {
                    "device_id": gpu_index,
                    "trt_engine_cache_enable": tensorrt_cache_path is not None,
                    "trt_engine_cache_path": tensorrt_cache_path,
                    "trt_fp16_enable": should_tensorrt_fp16,
                },
            ),
            (
                "CUDAExecutionProvider",
                {
                    "device_id": gpu_index,
                },
            ),
            "CPUExecutionProvider",
        ]
    elif execution_provider == "CUDAExecutionProvider":
        providers = [
            (
                "CUDAExecutionProvider",
                {
                    "device_id": gpu_index,
                },
            ),
            "CPUExecutionProvider",
        ]
    else:
        providers = [execution_provider, "CPUExecutionProvider"]

    session = ort.InferenceSession(model.bytes, providers=providers)
    return session


__session_cache: WeakKeyDictionary[
    any, ort.InferenceSession
] = WeakKeyDictionary()


def get_onnx_session(
    model: any,
    gpu_index: int,
    execution_provider: str,
    should_tensorrt_fp16: bool,
    tensorrt_cache_path: str | None = None,
) -> ort.InferenceSession:
    cached = __session_cache.get(model)
    if cached is None:
        cached = create_inference_session(
            model,
            gpu_index,
            execution_provider,
            should_tensorrt_fp16,
            tensorrt_cache_path,
        )
        __session_cache[model] = cached
    return cached

nodes/impl/onnx/session.py
vladmandic/sd-extension-chainner
from sys import float_info

import numpy as np
from onnx import numpy_helper as onph
from onnx.onnx_pb import AttributeProto, NodeProto, TensorProto
from nodes.log import logger

INT64_MIN, INT64_MAX = np.iinfo(np.int64).min, np.iinfo(np.int64).max
FLOAT32_MAX = float_info.max

APT = AttributeProto
TPT = TensorProto


def get_node_attr_ai(node: NodeProto, key: str) -> np.ndarray:
    for attr in node.attribute:
        if attr.name == key:
            return np.array(
                [max(min(i, INT64_MAX), INT64_MIN) for i in attr.ints], np.int64
            )

    return np.empty(0, np.int32)


def set_node_attr_ai(node: NodeProto, key: str, value: np.ndarray) -> None:
    attr_group = AttributeProto(name=key, floats=value, type=APT.INTS)
    node.attribute.append(attr_group)


def get_node_attr_af(node: NodeProto, key: str) -> np.ndarray:
    for attr in node.attribute:
        if attr.name == key:
            return np.array(list(attr.floats), np.float32)

    return np.empty(0, np.float32)


def get_node_attr_i(node: NodeProto, key: str, default: int = 0) -> int:
    for attr in node.attribute:
        if attr.name == key:
            return max(min(attr.i, INT64_MAX), INT64_MIN)

    return default


def get_node_attr_f(node: NodeProto, key: str, default: float = 0) -> float:
    for attr in node.attribute:
        if attr.name == key:
            return attr.f

    return default


def get_node_attr_s(node: NodeProto, key: str, default: str = ""):
    for attr in node.attribute:
        if attr.name == key:
            return attr.s.decode("ascii")

    return default


def get_node_attr_tensor(node: NodeProto, key: str) -> TensorProto:
    for attr in node.attribute:
        if attr.name == key:
            return attr.t

    return TensorProto()


def get_node_attr_from_input_f(tp: TensorProto) -> float:
    shape_data = onph.to_array(tp)

    if tp.data_type in (TPT.FLOAT, TPT.FLOAT16, TPT.DOUBLE, TPT.INT32):
        f = shape_data.item(0)
    elif tp.data_type == TPT.INT64:
        f = max(min(shape_data.item(0), INT64_MAX), INT64_MIN)
    else:
        raise TypeError(f"Unknown data type {tp.data_type}")

    return f


def get_node_attr_from_input_ai(tp: TensorProto) -> np.ndarray:
    if tp.data_type == TPT.INT32 or tp.data_type == TPT.INT64:
        shape_data = onph.to_array(tp)
        if shape_data.size == 1:
            shape_data = np.array([shape_data.item(0)], shape_data.dtype)
        return np.array(
            [
                max(min(val, INT64_MAX), INT64_MIN)
                if tp.data_type == TPT.INT64
                else val
                for val in shape_data
            ],
            shape_data.dtype,
        )
    else:
        logger.error(f"chaiNNer: unknown data type {tp.data_type}")

    return np.empty(0, np.int32)


def get_node_attr_from_input_af(tp: TensorProto) -> np.ndarray:
    if tp.data_type in (TPT.FLOAT, TPT.FLOAT16, TPT.DOUBLE):
        shape_data = onph.to_array(tp)
        return np.array(list(shape_data), shape_data.dtype)
    else:
        logger.error(f"chaiNNer: unknown data type {tp.data_type}")

    return np.empty(0, np.float32)


def get_tensor_proto_data_size(tp: TensorProto, fpmode: int = TPT.FLOAT) -> int:
    if tp.raw_data:
        if fpmode == TPT.FLOAT16:
            return len(tp.raw_data) // 2
        return len(tp.raw_data) // 4
    elif tp.data_type == TPT.FLOAT or tp.data_type == TPT.FLOAT16:
        return len(tp.float_data)

    return 0

nodes/impl/onnx/tensorproto_utils.py
vladmandic/sd-extension-chainner
from __future__ import annotations

from typing import Literal, Tuple, TYPE_CHECKING

import onnxoptimizer

if TYPE_CHECKING:
    import onnxruntime as ort
    from onnx.onnx_pb import ModelProto

OnnxInputShape = Literal["BCHW", "BHWC"]


def as_int(value) -> int | None:
    if isinstance(value, int):
        return value
    return None


def parse_onnx_shape(
    shape: Tuple[int | str, str | int, str | int, str | int]
) -> Tuple[OnnxInputShape, int, int | None, int | None]:
    if isinstance(shape[1], int) and shape[1] <= 4:
        return "BCHW", shape[1], as_int(shape[3]), as_int(shape[2])
    elif isinstance(shape[3], int) and shape[3] <= 4:
        return "BHWC", shape[3], as_int(shape[2]), as_int(shape[1])
    else:
        return "BCHW", 3, as_int(shape[3]), as_int(shape[2])


def get_input_shape(
    session: ort.InferenceSession,
) -> Tuple[OnnxInputShape, int, int | None, int | None]:
    """
    Returns the input shape, input channels, input width (optional), and input height (optional).
    """

    return parse_onnx_shape(session.get_inputs()[0].shape)


def get_output_shape(
    session: ort.InferenceSession,
) -> Tuple[OnnxInputShape, int, int | None, int | None]:
    """
    Returns the output shape, output channels, output width (optional), and output height (optional).
    """

    return parse_onnx_shape(session.get_outputs()[0].shape)


def safely_optimize_onnx_model(model_proto: ModelProto) -> ModelProto:
    """
    Optimizes the model using onnxoptimizer. If onnxoptimizer is not installed, the model is returned as is.
    """
    try:
        passes = onnxoptimizer.get_fuse_and_elimination_passes()
        model_proto = onnxoptimizer.optimize(model_proto, passes)
    except:
        pass
    return model_proto

nodes/impl/onnx/utils.py
vladmandic/sd-extension-chainner
from enum import Enum
from typing import Tuple

import numpy as np
from PIL import Image

from ..utils.utils import get_h_w_c
from .image_utils import FillColor, convert_to_BGRA, normalize, to_uint8


class InterpolationMethod(Enum):
    AUTO = -1
    NEAREST = 0
    LANCZOS = 1
    LINEAR = 2
    CUBIC = 3
    BOX = 4


class RotationInterpolationMethod(Enum):
    CUBIC = InterpolationMethod.CUBIC.value
    LINEAR = InterpolationMethod.LINEAR.value
    NEAREST = InterpolationMethod.NEAREST.value

    @property
    def interpolation_method(self) -> InterpolationMethod:
        return InterpolationMethod(self.value)


INTERPOLATION_METHODS_MAP = {
    InterpolationMethod.NEAREST: Image.NEAREST,
    InterpolationMethod.BOX: Image.BOX,
    InterpolationMethod.LINEAR: Image.BILINEAR,
    InterpolationMethod.CUBIC: Image.BICUBIC,
    InterpolationMethod.LANCZOS: Image.LANCZOS,
}


class RotateSizeChange(Enum):
    EXPAND = 1
    CROP = 0


def resize(
    img: np.ndarray, out_dims: Tuple[int, int], interpolation: InterpolationMethod
) -> np.ndarray:
    """Perform PIL resize"""

    if interpolation == InterpolationMethod.AUTO:
        # automatically chose a method that works
        new_w, new_h = out_dims
        old_h, old_w, _ = get_h_w_c(img)
        if new_w > old_w or new_h > old_h:
            interpolation = InterpolationMethod.LANCZOS
        else:
            interpolation = InterpolationMethod.BOX

    resample = INTERPOLATION_METHODS_MAP[interpolation]

    pimg = Image.fromarray(to_uint8(img, normalized=True))
    pimg = pimg.resize(out_dims, resample=resample)  # type: ignore
    return normalize(np.array(pimg))


def rotate(
    img: np.ndarray,
    angle: float,
    interpolation: RotationInterpolationMethod,
    expand: RotateSizeChange,
    fill: FillColor,
) -> np.ndarray:
    """Perform PIL rotate"""

    c = get_h_w_c(img)[2]
    if fill == FillColor.TRANSPARENT:
        img = convert_to_BGRA(img, c)
    fill_color = tuple([x * 255 for x in fill.get_color(c)])

    resample = INTERPOLATION_METHODS_MAP[interpolation.interpolation_method]

    pimg = Image.fromarray(to_uint8(img, normalized=True))
    pimg = pimg.rotate(
        angle,
        resample=resample,  # type: ignore
        expand=bool(expand.value),
        fillcolor=fill_color,
    )
    return normalize(np.array(pimg))

nodes/impl/pil_utils.py
vladmandic/sd-extension-chainner
# pylint: skip-file
import math
import re

import numpy as np
import torch
import torch.nn as nn
import torch.utils.checkpoint as checkpoint
from einops import rearrange
from einops.layers.torch import Rearrange
from torch import Tensor
from torch.nn import functional as F

from .timm.drop import DropPath
from .timm.weight_init import trunc_normal_


def img2windows(img, H_sp, W_sp):
    """
    Input: Image (B, C, H, W)
    Output: Window Partition (B', N, C)
    """
    B, C, H, W = img.shape
    img_reshape = img.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp)
    img_perm = (
        img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, H_sp * W_sp, C)
    )
    return img_perm


def windows2img(img_splits_hw, H_sp, W_sp, H, W):
    """
    Input: Window Partition (B', N, C)
    Output: Image (B, H, W, C)
    """
    B = int(img_splits_hw.shape[0] / (H * W / H_sp / W_sp))

    img = img_splits_hw.view(B, H // H_sp, W // W_sp, H_sp, W_sp, -1)
    img = img.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return img


class SpatialGate(nn.Module):
    """Spatial-Gate.
    Args:
        dim (int): Half of input channels.
    """

    def __init__(self, dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.conv = nn.Conv2d(
            dim, dim, kernel_size=3, stride=1, padding=1, groups=dim
        )  # DW Conv

    def forward(self, x, H, W):
        # Split
        x1, x2 = x.chunk(2, dim=-1)
        B, N, C = x.shape
        x2 = (
            self.conv(self.norm(x2).transpose(1, 2).contiguous().view(B, C // 2, H, W))
            .flatten(2)
            .transpose(-1, -2)
            .contiguous()
        )

        return x1 * x2


class SGFN(nn.Module):
    """Spatial-Gate Feed-Forward Network.
    Args:
        in_features (int): Number of input channels.
        hidden_features (int | None): Number of hidden channels. Default: None
        out_features (int | None): Number of output channels. Default: None
        act_layer (nn.Module): Activation layer. Default: nn.GELU
        drop (float): Dropout rate. Default: 0.0
    """

    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        drop=0.0,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.sg = SpatialGate(hidden_features // 2)
        self.fc2 = nn.Linear(hidden_features // 2, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x, H, W):
        """
        Input: x: (B, H*W, C), H, W
        Output: x: (B, H*W, C)
        """
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)

        x = self.sg(x, H, W)
        x = self.drop(x)

        x = self.fc2(x)
        x = self.drop(x)
        return x


class DynamicPosBias(nn.Module):
    # The implementation builds on Crossformer code https://github.com/cheerss/CrossFormer/blob/main/models/crossformer.py
    """Dynamic Relative Position Bias.
    Args:
        dim (int): Number of input channels.
        num_heads (int): Number of attention heads.
        residual (bool):  If True, use residual strage to connect conv.
    """

    def __init__(self, dim, num_heads, residual):
        super().__init__()
        self.residual = residual
        self.num_heads = num_heads
        self.pos_dim = dim // 4
        self.pos_proj = nn.Linear(2, self.pos_dim)
        self.pos1 = nn.Sequential(
            nn.LayerNorm(self.pos_dim),
            nn.ReLU(inplace=True),
            nn.Linear(self.pos_dim, self.pos_dim),
        )
        self.pos2 = nn.Sequential(
            nn.LayerNorm(self.pos_dim),
            nn.ReLU(inplace=True),
            nn.Linear(self.pos_dim, self.pos_dim),
        )
        self.pos3 = nn.Sequential(
            nn.LayerNorm(self.pos_dim),
            nn.ReLU(inplace=True),
            nn.Linear(self.pos_dim, self.num_heads),
        )

    def forward(self, biases):
        if self.residual:
            pos = self.pos_proj(biases)  # 2Gh-1 * 2Gw-1, heads
            pos = pos + self.pos1(pos)
            pos = pos + self.pos2(pos)
            pos = self.pos3(pos)
        else:
            pos = self.pos3(self.pos2(self.pos1(self.pos_proj(biases))))
        return pos


class Spatial_Attention(nn.Module):
    """Spatial Window Self-Attention.
    It supports rectangle window (containing square window).
    Args:
        dim (int): Number of input channels.
        idx (int): The indentix of window. (0/1)
        split_size (tuple(int)): Height and Width of spatial window.
        dim_out (int | None): The dimension of the attention output. Default: None
        num_heads (int): Number of attention heads. Default: 6
        attn_drop (float): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float): Dropout ratio of output. Default: 0.0
        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set
        position_bias (bool): The dynamic relative position bias. Default: True
    """

    def __init__(
        self,
        dim,
        idx,
        split_size=None,
        dim_out=None,
        num_heads=6,
        attn_drop=0.0,
        proj_drop=0.0,
        qk_scale=None,
        position_bias=True,
    ):
        if split_size is None:
            split_size = [8, 8]
        super().__init__()
        self.dim = dim
        self.dim_out = dim_out or dim
        self.split_size = split_size
        self.num_heads = num_heads
        self.idx = idx
        self.position_bias = position_bias

        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        if idx == 0:
            H_sp, W_sp = self.split_size[0], self.split_size[1]
        elif idx == 1:
            W_sp, H_sp = self.split_size[0], self.split_size[1]
        else:
            print("ERROR MODE", idx)
            exit(0)
        self.H_sp = H_sp
        self.W_sp = W_sp

        if self.position_bias:
            self.pos = DynamicPosBias(self.dim // 4, self.num_heads, residual=False)
            # generate mother-set
            position_bias_h = torch.arange(1 - self.H_sp, self.H_sp)
            position_bias_w = torch.arange(1 - self.W_sp, self.W_sp)
            biases = torch.stack(torch.meshgrid([position_bias_h, position_bias_w]))
            biases = biases.flatten(1).transpose(0, 1).contiguous().float()
            self.register_buffer("rpe_biases", biases)

            # get pair-wise relative position index for each token inside the window
            coords_h = torch.arange(self.H_sp)
            coords_w = torch.arange(self.W_sp)
            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))
            coords_flatten = torch.flatten(coords, 1)
            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
            relative_coords = relative_coords.permute(1, 2, 0).contiguous()
            relative_coords[:, :, 0] += self.H_sp - 1
            relative_coords[:, :, 1] += self.W_sp - 1
            relative_coords[:, :, 0] *= 2 * self.W_sp - 1
            relative_position_index = relative_coords.sum(-1)
            self.register_buffer("relative_position_index", relative_position_index)

        self.attn_drop = nn.Dropout(attn_drop)

    def im2win(self, x, H, W):
        B, N, C = x.shape
        x = x.transpose(-2, -1).contiguous().view(B, C, H, W)
        x = img2windows(x, self.H_sp, self.W_sp)
        x = (
            x.reshape(-1, self.H_sp * self.W_sp, self.num_heads, C // self.num_heads)
            .permute(0, 2, 1, 3)
            .contiguous()
        )
        return x

    def forward(self, qkv, H, W, mask=None):
        """
        Input: qkv: (B, 3*L, C), H, W, mask: (B, N, N), N is the window size
        Output: x (B, H, W, C)
        """
        q, k, v = qkv[0], qkv[1], qkv[2]

        B, L, C = q.shape
        assert L == H * W, "flatten img_tokens has wrong size"

        # partition the q,k,v, image to window
        q = self.im2win(q, H, W)
        k = self.im2win(k, H, W)
        v = self.im2win(v, H, W)

        q = q * self.scale
        attn = q @ k.transpose(-2, -1)  # B head N C @ B head C N --> B head N N

        # calculate drpe
        if self.position_bias:
            pos = self.pos(self.rpe_biases)
            # select position bias
            relative_position_bias = pos[self.relative_position_index.view(-1)].view(
                self.H_sp * self.W_sp, self.H_sp * self.W_sp, -1
            )
            relative_position_bias = relative_position_bias.permute(
                2, 0, 1
            ).contiguous()
            attn = attn + relative_position_bias.unsqueeze(0)

        N = attn.shape[3]

        # use mask for shift window
        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(
                0
            )
            attn = attn.view(-1, self.num_heads, N, N)

        attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype)
        attn = self.attn_drop(attn)

        x = attn @ v
        x = x.transpose(1, 2).reshape(
            -1, self.H_sp * self.W_sp, C
        )  # B head N N @ B head N C

        # merge the window, window to image
        x = windows2img(x, self.H_sp, self.W_sp, H, W)  # B H' W' C

        return x


class Adaptive_Spatial_Attention(nn.Module):
    # The implementation builds on CAT code https://github.com/Zhengchen1999/CAT
    """Adaptive Spatial Self-Attention
    Args:
        dim (int): Number of input channels.
        num_heads (int): Number of attention heads. Default: 6
        split_size (tuple(int)): Height and Width of spatial window.
        shift_size (tuple(int)): Shift size for spatial window.
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set.
        drop (float): Dropout rate. Default: 0.0
        attn_drop (float): Attention dropout rate. Default: 0.0
        rg_idx (int): The indentix of Residual Group (RG)
        b_idx (int): The indentix of Block in each RG
    """

    def __init__(
        self,
        dim,
        num_heads,
        reso=64,
        split_size=None,
        shift_size=None,
        qkv_bias=False,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        rg_idx=0,
        b_idx=0,
    ):
        if shift_size is None:
            shift_size = [1, 2]
        if split_size is None:
            split_size = [8, 8]
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.split_size = split_size
        self.shift_size = shift_size
        self.b_idx = b_idx
        self.rg_idx = rg_idx
        self.patches_resolution = reso
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)

        assert (
            0 <= self.shift_size[0] < self.split_size[0]
        ), "shift_size must in 0-split_size0"
        assert (
            0 <= self.shift_size[1] < self.split_size[1]
        ), "shift_size must in 0-split_size1"

        self.branch_num = 2

        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(drop)

        self.attns = nn.ModuleList(
            [
                Spatial_Attention(
                    dim // 2,
                    idx=i,
                    split_size=split_size,
                    num_heads=num_heads // 2,
                    dim_out=dim // 2,
                    qk_scale=qk_scale,
                    attn_drop=attn_drop,
                    proj_drop=drop,
                    position_bias=True,
                )
                for i in range(self.branch_num)
            ]
        )

        if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or (
            self.rg_idx % 2 != 0 and self.b_idx % 4 == 0
        ):
            attn_mask = self.calculate_mask(
                self.patches_resolution, self.patches_resolution
            )
            self.register_buffer("attn_mask_0", attn_mask[0])
            self.register_buffer("attn_mask_1", attn_mask[1])
        else:
            attn_mask = None
            self.register_buffer("attn_mask_0", None)
            self.register_buffer("attn_mask_1", None)

        self.dwconv = nn.Sequential(
            nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim),
            nn.BatchNorm2d(dim),
            nn.GELU(),
        )
        self.channel_interaction = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(dim, dim // 8, kernel_size=1),
            nn.BatchNorm2d(dim // 8),
            nn.GELU(),
            nn.Conv2d(dim // 8, dim, kernel_size=1),
        )
        self.spatial_interaction = nn.Sequential(
            nn.Conv2d(dim, dim // 16, kernel_size=1),
            nn.BatchNorm2d(dim // 16),
            nn.GELU(),
            nn.Conv2d(dim // 16, 1, kernel_size=1),
        )

    def calculate_mask(self, H, W):
        # The implementation builds on Swin Transformer code https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
        # calculate attention mask for shift window
        img_mask_0 = torch.zeros((1, H, W, 1))  # 1 H W 1 idx=0
        img_mask_1 = torch.zeros((1, H, W, 1))  # 1 H W 1 idx=1
        h_slices_0 = (
            slice(0, -self.split_size[0]),
            slice(-self.split_size[0], -self.shift_size[0]),
            slice(-self.shift_size[0], None),
        )
        w_slices_0 = (
            slice(0, -self.split_size[1]),
            slice(-self.split_size[1], -self.shift_size[1]),
            slice(-self.shift_size[1], None),
        )

        h_slices_1 = (
            slice(0, -self.split_size[1]),
            slice(-self.split_size[1], -self.shift_size[1]),
            slice(-self.shift_size[1], None),
        )
        w_slices_1 = (
            slice(0, -self.split_size[0]),
            slice(-self.split_size[0], -self.shift_size[0]),
            slice(-self.shift_size[0], None),
        )
        cnt = 0
        for h in h_slices_0:
            for w in w_slices_0:
                img_mask_0[:, h, w, :] = cnt
                cnt += 1
        cnt = 0
        for h in h_slices_1:
            for w in w_slices_1:
                img_mask_1[:, h, w, :] = cnt
                cnt += 1

        # calculate mask for window-0
        img_mask_0 = img_mask_0.view(
            1,
            H // self.split_size[0],
            self.split_size[0],
            W // self.split_size[1],
            self.split_size[1],
            1,
        )
        img_mask_0 = (
            img_mask_0.permute(0, 1, 3, 2, 4, 5)
            .contiguous()
            .view(-1, self.split_size[0], self.split_size[1], 1)
        )  # nW, sw[0], sw[1], 1
        mask_windows_0 = img_mask_0.view(-1, self.split_size[0] * self.split_size[1])
        attn_mask_0 = mask_windows_0.unsqueeze(1) - mask_windows_0.unsqueeze(2)
        attn_mask_0 = attn_mask_0.masked_fill(
            attn_mask_0 != 0, float(-100.0)
        ).masked_fill(attn_mask_0 == 0, float(0.0))

        # calculate mask for window-1
        img_mask_1 = img_mask_1.view(
            1,
            H // self.split_size[1],
            self.split_size[1],
            W // self.split_size[0],
            self.split_size[0],
            1,
        )
        img_mask_1 = (
            img_mask_1.permute(0, 1, 3, 2, 4, 5)
            .contiguous()
            .view(-1, self.split_size[1], self.split_size[0], 1)
        )  # nW, sw[1], sw[0], 1
        mask_windows_1 = img_mask_1.view(-1, self.split_size[1] * self.split_size[0])
        attn_mask_1 = mask_windows_1.unsqueeze(1) - mask_windows_1.unsqueeze(2)
        attn_mask_1 = attn_mask_1.masked_fill(
            attn_mask_1 != 0, float(-100.0)
        ).masked_fill(attn_mask_1 == 0, float(0.0))

        return attn_mask_0, attn_mask_1

    def forward(self, x, H, W):
        """
        Input: x: (B, H*W, C), H, W
        Output: x: (B, H*W, C)
        """
        B, L, C = x.shape
        assert L == H * W, "flatten img_tokens has wrong size"

        qkv = self.qkv(x).reshape(B, -1, 3, C).permute(2, 0, 1, 3)  # 3, B, HW, C
        # V without partition
        v = qkv[2].transpose(-2, -1).contiguous().view(B, C, H, W)

        # image padding
        max_split_size = max(self.split_size[0], self.split_size[1])
        pad_l = pad_t = 0
        pad_r = (max_split_size - W % max_split_size) % max_split_size
        pad_b = (max_split_size - H % max_split_size) % max_split_size

        qkv = qkv.reshape(3 * B, H, W, C).permute(0, 3, 1, 2)  # 3B C H W
        qkv = (
            F.pad(qkv, (pad_l, pad_r, pad_t, pad_b))
            .reshape(3, B, C, -1)
            .transpose(-2, -1)
        )  # l r t b
        _H = pad_b + H
        _W = pad_r + W
        _L = _H * _W

        # window-0 and window-1 on split channels [C/2, C/2]; for square windows (e.g., 8x8), window-0 and window-1 can be merged
        # shift in block: (0, 4, 8, ...), (2, 6, 10, ...), (0, 4, 8, ...), (2, 6, 10, ...), ...
        if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or (
            self.rg_idx % 2 != 0 and self.b_idx % 4 == 0
        ):
            qkv = qkv.view(3, B, _H, _W, C)
            qkv_0 = torch.roll(
                qkv[:, :, :, :, : C // 2],
                shifts=(-self.shift_size[0], -self.shift_size[1]),
                dims=(2, 3),
            )
            qkv_0 = qkv_0.view(3, B, _L, C // 2)
            qkv_1 = torch.roll(
                qkv[:, :, :, :, C // 2 :],
                shifts=(-self.shift_size[1], -self.shift_size[0]),
                dims=(2, 3),
            )
            qkv_1 = qkv_1.view(3, B, _L, C // 2)

            if self.patches_resolution != _H or self.patches_resolution != _W:
                mask_tmp = self.calculate_mask(_H, _W)
                x1_shift = self.attns[0](qkv_0, _H, _W, mask=mask_tmp[0].to(x.device))
                x2_shift = self.attns[1](qkv_1, _H, _W, mask=mask_tmp[1].to(x.device))
            else:
                x1_shift = self.attns[0](qkv_0, _H, _W, mask=self.attn_mask_0)
                x2_shift = self.attns[1](qkv_1, _H, _W, mask=self.attn_mask_1)

            x1 = torch.roll(
                x1_shift, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2)
            )
            x2 = torch.roll(
                x2_shift, shifts=(self.shift_size[1], self.shift_size[0]), dims=(1, 2)
            )
            x1 = x1[:, :H, :W, :].reshape(B, L, C // 2)
            x2 = x2[:, :H, :W, :].reshape(B, L, C // 2)
            # attention output
            attened_x = torch.cat([x1, x2], dim=2)

        else:
            x1 = self.attns[0](qkv[:, :, :, : C // 2], _H, _W)[:, :H, :W, :].reshape(
                B, L, C // 2
            )
            x2 = self.attns[1](qkv[:, :, :, C // 2 :], _H, _W)[:, :H, :W, :].reshape(
                B, L, C // 2
            )
            # attention output
            attened_x = torch.cat([x1, x2], dim=2)

        # convolution output
        conv_x = self.dwconv(v)

        # Adaptive Interaction Module (AIM)
        # C-Map (before sigmoid)
        channel_map = (
            self.channel_interaction(conv_x)
            .permute(0, 2, 3, 1)
            .contiguous()
            .view(B, 1, C)
        )
        # S-Map (before sigmoid)
        attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W)
        spatial_map = self.spatial_interaction(attention_reshape)

        # C-I
        attened_x = attened_x * torch.sigmoid(channel_map)
        # S-I
        conv_x = torch.sigmoid(spatial_map) * conv_x
        conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(B, L, C)

        x = attened_x + conv_x

        x = self.proj(x)
        x = self.proj_drop(x)

        return x


class Adaptive_Channel_Attention(nn.Module):
    # The implementation builds on XCiT code https://github.com/facebookresearch/xcit
    """Adaptive Channel Self-Attention
    Args:
        dim (int): Number of input channels.
        num_heads (int): Number of attention heads. Default: 6
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set.
        attn_drop (float): Attention dropout rate. Default: 0.0
        drop_path (float): Stochastic depth rate. Default: 0.0
    """

    def __init__(
        self,
        dim,
        num_heads=8,
        qkv_bias=False,
        qk_scale=None,
        attn_drop=0.0,
        proj_drop=0.0,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.dwconv = nn.Sequential(
            nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim),
            nn.BatchNorm2d(dim),
            nn.GELU(),
        )
        self.channel_interaction = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(dim, dim // 8, kernel_size=1),
            nn.BatchNorm2d(dim // 8),
            nn.GELU(),
            nn.Conv2d(dim // 8, dim, kernel_size=1),
        )
        self.spatial_interaction = nn.Sequential(
            nn.Conv2d(dim, dim // 16, kernel_size=1),
            nn.BatchNorm2d(dim // 16),
            nn.GELU(),
            nn.Conv2d(dim // 16, 1, kernel_size=1),
        )

    def forward(self, x, H, W):
        """
        Input: x: (B, H*W, C), H, W
        Output: x: (B, H*W, C)
        """
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        q = q.transpose(-2, -1)
        k = k.transpose(-2, -1)
        v = v.transpose(-2, -1)

        v_ = v.reshape(B, C, N).contiguous().view(B, C, H, W)

        q = torch.nn.functional.normalize(q, dim=-1)
        k = torch.nn.functional.normalize(k, dim=-1)

        attn = (q @ k.transpose(-2, -1)) * self.temperature
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        # attention output
        attened_x = (attn @ v).permute(0, 3, 1, 2).reshape(B, N, C)

        # convolution output
        conv_x = self.dwconv(v_)

        # Adaptive Interaction Module (AIM)
        # C-Map (before sigmoid)
        attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W)
        channel_map = self.channel_interaction(attention_reshape)
        # S-Map (before sigmoid)
        spatial_map = (
            self.spatial_interaction(conv_x)
            .permute(0, 2, 3, 1)
            .contiguous()
            .view(B, N, 1)
        )

        # S-I
        attened_x = attened_x * torch.sigmoid(spatial_map)
        # C-I
        conv_x = conv_x * torch.sigmoid(channel_map)
        conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(B, N, C)

        x = attened_x + conv_x

        x = self.proj(x)
        x = self.proj_drop(x)

        return x


class DATB(nn.Module):
    def __init__(
        self,
        dim,
        num_heads,
        reso=64,
        split_size=None,
        shift_size=None,
        expansion_factor=4.0,
        qkv_bias=False,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        rg_idx=0,
        b_idx=0,
    ):
        if shift_size is None:
            shift_size = [1, 2]
        if split_size is None:
            split_size = [2, 4]
        super().__init__()

        self.norm1 = norm_layer(dim)

        if b_idx % 2 == 0:
            # DSTB
            self.attn = Adaptive_Spatial_Attention(
                dim,
                num_heads=num_heads,
                reso=reso,
                split_size=split_size,
                shift_size=shift_size,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop=drop,
                attn_drop=attn_drop,
                rg_idx=rg_idx,
                b_idx=b_idx,
            )
        else:
            # DCTB
            self.attn = Adaptive_Channel_Attention(
                dim,
                num_heads=num_heads,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                attn_drop=attn_drop,
                proj_drop=drop,
            )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

        ffn_hidden_dim = int(dim * expansion_factor)
        self.ffn = SGFN(
            in_features=dim,
            hidden_features=ffn_hidden_dim,
            out_features=dim,
            act_layer=act_layer,
        )
        self.norm2 = norm_layer(dim)

    def forward(self, x, x_size):
        """
        Input: x: (B, H*W, C), x_size: (H, W)
        Output: x: (B, H*W, C)
        """
        H, W = x_size
        x = x + self.drop_path(self.attn(self.norm1(x), H, W))
        x = x + self.drop_path(self.ffn(self.norm2(x), H, W))

        return x


class ResidualGroup(nn.Module):
    """ResidualGroup
    Args:
        dim (int): Number of input channels.
        reso (int): Input resolution.
        num_heads (int): Number of attention heads.
        split_size (tuple(int)): Height and Width of spatial window.
        expansion_factor (float): Ratio of ffn hidden dim to embedding dim.
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. Default: None
        drop (float): Dropout rate. Default: 0
        attn_drop(float): Attention dropout rate. Default: 0
        drop_paths (float | None): Stochastic depth rate.
        act_layer (nn.Module): Activation layer. Default: nn.GELU
        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm
        depth (int): Number of dual aggregation Transformer blocks in residual group.
        use_chk (bool): Whether to use checkpointing to save memory.
        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
    """

    def __init__(
        self,
        dim,
        reso,
        num_heads,
        split_size=None,
        expansion_factor=4.0,
        qkv_bias=False,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_paths=None,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        depth=2,
        use_chk=False,
        resi_connection="1conv",
        rg_idx=0,
    ):
        if split_size is None:
            split_size = [2, 4]
        super().__init__()
        self.use_chk = use_chk
        self.reso = reso

        self.blocks = nn.ModuleList(
            [
                DATB(
                    dim=dim,
                    num_heads=num_heads,
                    reso=reso,
                    split_size=split_size,
                    shift_size=[split_size[0] // 2, split_size[1] // 2],
                    expansion_factor=expansion_factor,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop,
                    attn_drop=attn_drop,
                    drop_path=drop_paths[i],
                    act_layer=act_layer,
                    norm_layer=norm_layer,
                    rg_idx=rg_idx,
                    b_idx=i,
                )
                for i in range(depth)
            ]
        )

        if resi_connection == "1conv":
            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
        elif resi_connection == "3conv":
            self.conv = nn.Sequential(
                nn.Conv2d(dim, dim // 4, 3, 1, 1),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(dim // 4, dim, 3, 1, 1),
            )

    def forward(self, x, x_size):
        """
        Input: x: (B, H*W, C), x_size: (H, W)
        Output: x: (B, H*W, C)
        """
        H, W = x_size
        res = x
        for blk in self.blocks:
            if self.use_chk:
                x = checkpoint.checkpoint(blk, x, x_size)
            else:
                x = blk(x, x_size)
        x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W)
        x = self.conv(x)
        x = rearrange(x, "b c h w -> b (h w) c")
        x = res + x

        return x


class Upsample(nn.Sequential):
    """Upsample module.
    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.
    """

    def __init__(self, scale, num_feat):
        m = []
        if (scale & (scale - 1)) == 0:  # scale = 2^n
            for _ in range(int(math.log(scale, 2))):
                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
                m.append(nn.PixelShuffle(2))
        elif scale == 3:
            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
            m.append(nn.PixelShuffle(3))
        else:
            raise ValueError(
                f"scale {scale} is not supported. " "Supported scales: 2^n and 3."
            )
        super(Upsample, self).__init__(*m)


class UpsampleOneStep(nn.Sequential):
    """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
       Used in lightweight SR to save parameters.

    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.

    """

    def __init__(self, scale, num_feat, num_out_ch, input_resolution=None):
        self.num_feat = num_feat
        self.input_resolution = input_resolution
        m = []
        m.append(nn.Conv2d(num_feat, (scale**2) * num_out_ch, 3, 1, 1))
        m.append(nn.PixelShuffle(scale))
        super(UpsampleOneStep, self).__init__(*m)

    def flops(self):
        h, w = self.input_resolution
        flops = h * w * self.num_feat * 3 * 9
        return flops


class DAT(nn.Module):
    """Dual Aggregation Transformer
    Args:
        img_size (int): Input image size. Default: 64
        in_chans (int): Number of input image channels. Default: 3
        embed_dim (int): Patch embedding dimension. Default: 180
        depths (tuple(int)): Depth of each residual group (number of DATB in each RG).
        split_size (tuple(int)): Height and Width of spatial window.
        num_heads (tuple(int)): Number of attention heads in different residual groups.
        expansion_factor (float): Ratio of ffn hidden dim to embedding dim. Default: 4
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. Default: None
        drop_rate (float): Dropout rate. Default: 0
        attn_drop_rate (float): Attention dropout rate. Default: 0
        drop_path_rate (float): Stochastic depth rate. Default: 0.1
        act_layer (nn.Module): Activation layer. Default: nn.GELU
        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm
        use_chk (bool): Whether to use checkpointing to save memory.
        upscale: Upscale factor. 2/3/4 for image SR
        img_range: Image range. 1. or 255.
        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
    """

    def __init__(self, state_dict):
        super().__init__()

        # defaults
        img_size = 64
        in_chans = 3
        embed_dim = 180
        split_size = [2, 4]
        depth = [2, 2, 2, 2]
        num_heads = [2, 2, 2, 2]
        expansion_factor = 4.0
        qkv_bias = True
        qk_scale = None
        drop_rate = 0.0
        attn_drop_rate = 0.0
        drop_path_rate = 0.1
        act_layer = nn.GELU
        norm_layer = nn.LayerNorm
        use_chk = False
        upscale = 2
        img_range = 1.0
        resi_connection = "1conv"
        upsampler = "pixelshuffle"

        self.model_arch = "DAT"
        self.sub_type = "SR"
        self.state = state_dict

        state_keys = state_dict.keys()
        if "conv_before_upsample.0.weight" in state_keys:
            if "conv_up1.weight" in state_keys:
                upsampler = "nearest+conv"
            else:
                upsampler = "pixelshuffle"
        elif "upsample.0.weight" in state_keys:
            upsampler = "pixelshuffledirect"
        else:
            upsampler = ""

        num_feat = (
            state_dict.get("conv_before_upsample.0.weight", None).shape[1]
            if state_dict.get("conv_before_upsample.weight", None)
            else 64
        )

        num_in_ch = state_dict["conv_first.weight"].shape[1]
        in_chans = num_in_ch
        if "conv_last.weight" in state_keys:
            num_out_ch = state_dict["conv_last.weight"].shape[0]
        else:
            num_out_ch = num_in_ch

        upscale = 1
        if upsampler == "nearest+conv":
            upsample_keys = [
                x for x in state_keys if "conv_up" in x and "bias" not in x
            ]

            for upsample_key in upsample_keys:
                upscale *= 2
        elif upsampler == "pixelshuffle":
            upsample_keys = [
                x
                for x in state_keys
                if "upsample" in x and "conv" not in x and "bias" not in x
            ]
            for upsample_key in upsample_keys:
                shape = state_dict[upsample_key].shape[0]
                upscale *= math.sqrt(shape // num_feat)
            upscale = int(upscale)
        elif upsampler == "pixelshuffledirect":
            upscale = int(
                math.sqrt(state_dict["upsample.0.bias"].shape[0] // num_out_ch)
            )

        max_layer_num = 0
        max_block_num = 0
        for key in state_keys:
            result = re.match(r"layers.(\d*).blocks.(\d*).norm1.weight", key)
            if result:
                layer_num, block_num = result.groups()
                max_layer_num = max(max_layer_num, int(layer_num))
                max_block_num = max(max_block_num, int(block_num))

        depth = [max_block_num + 1 for _ in range(max_layer_num + 1)]

        if "layers.0.blocks.1.attn.temperature" in state_keys:
            num_heads_num = state_dict["layers.0.blocks.1.attn.temperature"].shape[0]
            num_heads = [num_heads_num for _ in range(max_layer_num + 1)]
        else:
            num_heads = depth

        embed_dim = state_dict["conv_first.weight"].shape[0]
        expansion_factor = float(
            state_dict["layers.0.blocks.0.ffn.fc1.weight"].shape[0] / embed_dim
        )

        # TODO: could actually count the layers, but this should do
        if "layers.0.conv.4.weight" in state_keys:
            resi_connection = "3conv"
        else:
            resi_connection = "1conv"

        if "layers.0.blocks.2.attn.attn_mask_0" in state_keys:
            attn_mask_0_x, attn_mask_0_y, attn_mask_0_z = state_dict[
                "layers.0.blocks.2.attn.attn_mask_0"
            ].shape

            img_size = int(math.sqrt(attn_mask_0_x * attn_mask_0_y))

        if "layers.0.blocks.0.attn.attns.0.rpe_biases" in state_keys:
            split_sizes = (
                state_dict["layers.0.blocks.0.attn.attns.0.rpe_biases"][-1] + 1
            )
            split_size = [int(x) for x in split_sizes]

        self.in_nc = num_in_ch
        self.out_nc = num_out_ch
        self.num_feat = num_feat
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.depth = depth
        self.scale = upscale
        self.upsampler = upsampler
        self.img_size = img_size
        self.img_range = img_range
        self.expansion_factor = expansion_factor
        self.resi_connection = resi_connection
        self.split_size = split_size

        self.supports_fp16 = False  # Too much weirdness to support this at the moment
        self.supports_bfp16 = True
        self.min_size_restriction = 16

        num_in_ch = in_chans
        num_out_ch = in_chans
        num_feat = 64
        self.img_range = img_range
        if in_chans == 3:
            rgb_mean = (0.4488, 0.4371, 0.4040)
            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
        else:
            self.mean = torch.zeros(1, 1, 1, 1)
        self.upscale = upscale
        self.upsampler = upsampler

        # ------------------------- 1, Shallow Feature Extraction ------------------------- #
        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)

        # ------------------------- 2, Deep Feature Extraction ------------------------- #
        self.num_layers = len(depth)
        self.use_chk = use_chk
        self.num_features = (
            self.embed_dim
        ) = embed_dim  # num_features for consistency with other models
        heads = num_heads

        self.before_RG = nn.Sequential(
            Rearrange("b c h w -> b (h w) c"), nn.LayerNorm(embed_dim)
        )

        curr_dim = embed_dim
        dpr = [
            x.item() for x in torch.linspace(0, drop_path_rate, np.sum(depth))
        ]  # stochastic depth decay rule

        self.layers = nn.ModuleList()
        for i in range(self.num_layers):
            layer = ResidualGroup(
                dim=embed_dim,
                num_heads=heads[i],
                reso=img_size,
                split_size=split_size,
                expansion_factor=expansion_factor,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop=drop_rate,
                attn_drop=attn_drop_rate,
                drop_paths=dpr[sum(depth[:i]) : sum(depth[: i + 1])],
                act_layer=act_layer,
                norm_layer=norm_layer,
                depth=depth[i],
                use_chk=use_chk,
                resi_connection=resi_connection,
                rg_idx=i,
            )
            self.layers.append(layer)

        self.norm = norm_layer(curr_dim)
        # build the last conv layer in deep feature extraction
        if resi_connection == "1conv":
            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
        elif resi_connection == "3conv":
            # to save parameters and memory
            self.conv_after_body = nn.Sequential(
                nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1),
            )

        # ------------------------- 3, Reconstruction ------------------------- #
        if self.upsampler == "pixelshuffle":
            # for classical SR
            self.conv_before_upsample = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.upsample = Upsample(upscale, num_feat)
            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
        elif self.upsampler == "pixelshuffledirect":
            # for lightweight SR (to save parameters)
            self.upsample = UpsampleOneStep(
                upscale, embed_dim, num_out_ch, (img_size, img_size)
            )

        self.apply(self._init_weights)
        self.load_state_dict(state_dict, strict=True)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(
            m, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d)
        ):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward_features(self, x):
        _, _, H, W = x.shape
        x_size = [H, W]
        x = self.before_RG(x)
        for layer in self.layers:
            x = layer(x, x_size)
        x = self.norm(x)
        x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W)

        return x

    def forward(self, x):
        """
        Input: x: (B, C, H, W)
        """
        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == "pixelshuffle":
            # for image SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == "pixelshuffledirect":
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)

        x = x / self.img_range + self.mean
        return x

nodes/impl/pytorch/architecture/DAT.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# HAT from https://github.com/XPixelGroup/HAT/blob/main/hat/archs/hat_arch.py
import math
import re

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

from .timm.helpers import to_2tuple
from .timm.weight_init import trunc_normal_


def drop_path(x, drop_prob: float = 0.0, training: bool = False):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
    From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py
    """
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (
        x.ndim - 1
    )  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py
    """

    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)  # type: ignore


class ChannelAttention(nn.Module):
    """Channel attention used in RCAN.
    Args:
        num_feat (int): Channel number of intermediate features.
        squeeze_factor (int): Channel squeeze factor. Default: 16.
    """

    def __init__(self, num_feat, squeeze_factor=16):
        super(ChannelAttention, self).__init__()
        self.attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(num_feat, num_feat // squeeze_factor, 1, padding=0),
            nn.ReLU(inplace=True),
            nn.Conv2d(num_feat // squeeze_factor, num_feat, 1, padding=0),
            nn.Sigmoid(),
        )

    def forward(self, x):
        y = self.attention(x)
        return x * y


class CAB(nn.Module):
    def __init__(self, num_feat, compress_ratio=3, squeeze_factor=30):
        super(CAB, self).__init__()

        self.cab = nn.Sequential(
            nn.Conv2d(num_feat, num_feat // compress_ratio, 3, 1, 1),
            nn.GELU(),
            nn.Conv2d(num_feat // compress_ratio, num_feat, 3, 1, 1),
            ChannelAttention(num_feat, squeeze_factor),
        )

    def forward(self, x):
        return self.cab(x)


class Mlp(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        drop=0.0,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


def window_partition(x, window_size):
    """
    Args:
        x: (b, h, w, c)
        window_size (int): window size
    Returns:
        windows: (num_windows*b, window_size, window_size, c)
    """
    b, h, w, c = x.shape
    x = x.view(b, h // window_size, window_size, w // window_size, window_size, c)
    windows = (
        x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, c)
    )
    return windows


def window_reverse(windows, window_size, h, w):
    """
    Args:
        windows: (num_windows*b, window_size, window_size, c)
        window_size (int): Window size
        h (int): Height of image
        w (int): Width of image
    Returns:
        x: (b, h, w, c)
    """
    b = int(windows.shape[0] / (h * w / window_size / window_size))
    x = windows.view(
        b, h // window_size, w // window_size, window_size, window_size, -1
    )
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(b, h, w, -1)
    return x


class WindowAttention(nn.Module):
    r"""Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.
    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
    """

    def __init__(
        self,
        dim,
        window_size,
        num_heads,
        qkv_bias=True,
        qk_scale=None,
        attn_drop=0.0,
        proj_drop=0.0,
    ):
        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wh, Ww
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        # define a parameter table of relative position bias
        self.relative_position_bias_table = nn.Parameter(  # type: ignore
            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
        )  # 2*Wh-1 * 2*Ww-1, nH

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)

        self.proj_drop = nn.Dropout(proj_drop)

        trunc_normal_(self.relative_position_bias_table, std=0.02)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x, rpi, mask=None):
        """
        Args:
            x: input features with shape of (num_windows*b, n, c)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        b_, n, c = x.shape
        qkv = (
            self.qkv(x)
            .reshape(b_, n, 3, self.num_heads, c // self.num_heads)
            .permute(2, 0, 3, 1, 4)
        )
        q, k, v = (
            qkv[0],
            qkv[1],
            qkv[2],
        )  # make torchscript happy (cannot use tensor as tuple)

        q = q * self.scale
        attn = q @ k.transpose(-2, -1)

        relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(
            self.window_size[0] * self.window_size[1],
            self.window_size[0] * self.window_size[1],
            -1,
        )  # Wh*Ww,Wh*Ww,nH
        relative_position_bias = relative_position_bias.permute(
            2, 0, 1
        ).contiguous()  # nH, Wh*Ww, Wh*Ww
        attn = attn + relative_position_bias.unsqueeze(0)

        if mask is not None:
            nw = mask.shape[0]
            attn = attn.view(b_ // nw, nw, self.num_heads, n, n) + mask.unsqueeze(
                1
            ).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, n, n)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(b_, n, c)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class HAB(nn.Module):
    r"""Hybrid Attention Block.
    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        num_heads (int): Number of attention heads.
        window_size (int): Window size.
        shift_size (int): Shift size for SW-MSA.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float, optional): Stochastic depth rate. Default: 0.0
        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(
        self,
        dim,
        input_resolution,
        num_heads,
        window_size=7,
        shift_size=0,
        compress_ratio=3,
        squeeze_factor=30,
        conv_scale=0.01,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        if min(self.input_resolution) <= self.window_size:
            # if window size is larger than input resolution, we don't partition windows
            self.shift_size = 0
            self.window_size = min(self.input_resolution)
        assert (
            0 <= self.shift_size < self.window_size
        ), "shift_size must in 0-window_size"

        self.norm1 = norm_layer(dim)
        self.attn = WindowAttention(
            dim,
            window_size=to_2tuple(self.window_size),
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            attn_drop=attn_drop,
            proj_drop=drop,
        )

        self.conv_scale = conv_scale
        self.conv_block = CAB(
            num_feat=dim, compress_ratio=compress_ratio, squeeze_factor=squeeze_factor
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

    def forward(self, x, x_size, rpi_sa, attn_mask):
        h, w = x_size
        b, _, c = x.shape
        # assert seq_len == h * w, "input feature has wrong size"

        shortcut = x
        x = self.norm1(x)
        x = x.view(b, h, w, c)

        # Conv_X
        conv_x = self.conv_block(x.permute(0, 3, 1, 2))
        conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(b, h * w, c)

        # cyclic shift
        if self.shift_size > 0:
            shifted_x = torch.roll(
                x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)
            )
            attn_mask = attn_mask
        else:
            shifted_x = x
            attn_mask = None

        # partition windows
        x_windows = window_partition(
            shifted_x, self.window_size
        )  # nw*b, window_size, window_size, c
        x_windows = x_windows.view(
            -1, self.window_size * self.window_size, c
        )  # nw*b, window_size*window_size, c

        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
        attn_windows = self.attn(x_windows, rpi=rpi_sa, mask=attn_mask)

        # merge windows
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, c)
        shifted_x = window_reverse(attn_windows, self.window_size, h, w)  # b h' w' c

        # reverse cyclic shift
        if self.shift_size > 0:
            attn_x = torch.roll(
                shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)
            )
        else:
            attn_x = shifted_x
        attn_x = attn_x.view(b, h * w, c)

        # FFN
        x = shortcut + self.drop_path(attn_x) + conv_x * self.conv_scale
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x


class PatchMerging(nn.Module):
    r"""Patch Merging Layer.
    Args:
        input_resolution (tuple[int]): Resolution of input feature.
        dim (int): Number of input channels.
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(4 * dim)

    def forward(self, x):
        """
        x: b, h*w, c
        """
        h, w = self.input_resolution
        b, seq_len, c = x.shape
        assert seq_len == h * w, "input feature has wrong size"
        assert h % 2 == 0 and w % 2 == 0, f"x size ({h}*{w}) are not even."

        x = x.view(b, h, w, c)

        x0 = x[:, 0::2, 0::2, :]  # b h/2 w/2 c
        x1 = x[:, 1::2, 0::2, :]  # b h/2 w/2 c
        x2 = x[:, 0::2, 1::2, :]  # b h/2 w/2 c
        x3 = x[:, 1::2, 1::2, :]  # b h/2 w/2 c
        x = torch.cat([x0, x1, x2, x3], -1)  # b h/2 w/2 4*c
        x = x.view(b, -1, 4 * c)  # b h/2*w/2 4*c

        x = self.norm(x)
        x = self.reduction(x)

        return x


class OCAB(nn.Module):
    # overlapping cross-attention block

    def __init__(
        self,
        dim,
        input_resolution,
        window_size,
        overlap_ratio,
        num_heads,
        qkv_bias=True,
        qk_scale=None,
        mlp_ratio=2,
        norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.window_size = window_size
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5
        self.overlap_win_size = int(window_size * overlap_ratio) + window_size

        self.norm1 = norm_layer(dim)
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.unfold = nn.Unfold(
            kernel_size=(self.overlap_win_size, self.overlap_win_size),
            stride=window_size,
            padding=(self.overlap_win_size - window_size) // 2,
        )

        # define a parameter table of relative position bias
        self.relative_position_bias_table = nn.Parameter(  # type: ignore
            torch.zeros(
                (window_size + self.overlap_win_size - 1)
                * (window_size + self.overlap_win_size - 1),
                num_heads,
            )
        )  # 2*Wh-1 * 2*Ww-1, nH

        trunc_normal_(self.relative_position_bias_table, std=0.02)
        self.softmax = nn.Softmax(dim=-1)

        self.proj = nn.Linear(dim, dim)

        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim, hidden_features=mlp_hidden_dim, act_layer=nn.GELU
        )

    def forward(self, x, x_size, rpi):
        h, w = x_size
        b, _, c = x.shape

        shortcut = x
        x = self.norm1(x)
        x = x.view(b, h, w, c)

        qkv = self.qkv(x).reshape(b, h, w, 3, c).permute(3, 0, 4, 1, 2)  # 3, b, c, h, w
        q = qkv[0].permute(0, 2, 3, 1)  # b, h, w, c
        kv = torch.cat((qkv[1], qkv[2]), dim=1)  # b, 2*c, h, w

        # partition windows
        q_windows = window_partition(
            q, self.window_size
        )  # nw*b, window_size, window_size, c
        q_windows = q_windows.view(
            -1, self.window_size * self.window_size, c
        )  # nw*b, window_size*window_size, c

        kv_windows = self.unfold(kv)  # b, c*w*w, nw
        kv_windows = rearrange(
            kv_windows,
            "b (nc ch owh oww) nw -> nc (b nw) (owh oww) ch",
            nc=2,
            ch=c,
            owh=self.overlap_win_size,
            oww=self.overlap_win_size,
        ).contiguous()  # 2, nw*b, ow*ow, c
        # Do the above rearrangement without the rearrange function
        # kv_windows = kv_windows.view(
        #     2, b, self.overlap_win_size, self.overlap_win_size, c, -1
        # )
        # kv_windows = kv_windows.permute(0, 5, 1, 2, 3, 4).contiguous()
        # kv_windows = kv_windows.view(
        #     2, -1, self.overlap_win_size * self.overlap_win_size, c
        # )

        k_windows, v_windows = kv_windows[0], kv_windows[1]  # nw*b, ow*ow, c

        b_, nq, _ = q_windows.shape
        _, n, _ = k_windows.shape
        d = self.dim // self.num_heads
        q = q_windows.reshape(b_, nq, self.num_heads, d).permute(
            0, 2, 1, 3
        )  # nw*b, nH, nq, d
        k = k_windows.reshape(b_, n, self.num_heads, d).permute(
            0, 2, 1, 3
        )  # nw*b, nH, n, d
        v = v_windows.reshape(b_, n, self.num_heads, d).permute(
            0, 2, 1, 3
        )  # nw*b, nH, n, d

        q = q * self.scale
        attn = q @ k.transpose(-2, -1)

        relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(
            self.window_size * self.window_size,
            self.overlap_win_size * self.overlap_win_size,
            -1,
        )  # ws*ws, wse*wse, nH
        relative_position_bias = relative_position_bias.permute(
            2, 0, 1
        ).contiguous()  # nH, ws*ws, wse*wse
        attn = attn + relative_position_bias.unsqueeze(0)

        attn = self.softmax(attn)
        attn_windows = (attn @ v).transpose(1, 2).reshape(b_, nq, self.dim)

        # merge windows
        attn_windows = attn_windows.view(
            -1, self.window_size, self.window_size, self.dim
        )
        x = window_reverse(attn_windows, self.window_size, h, w)  # b h w c
        x = x.view(b, h * w, self.dim)

        x = self.proj(x) + shortcut

        x = x + self.mlp(self.norm2(x))
        return x


class AttenBlocks(nn.Module):
    """A series of attention blocks for one RHAG.
    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
    """

    def __init__(
        self,
        dim,
        input_resolution,
        depth,
        num_heads,
        window_size,
        compress_ratio,
        squeeze_factor,
        conv_scale,
        overlap_ratio,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        downsample=None,
        use_checkpoint=False,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth
        self.use_checkpoint = use_checkpoint

        # build blocks
        self.blocks = nn.ModuleList(
            [
                HAB(
                    dim=dim,
                    input_resolution=input_resolution,
                    num_heads=num_heads,
                    window_size=window_size,
                    shift_size=0 if (i % 2 == 0) else window_size // 2,
                    compress_ratio=compress_ratio,
                    squeeze_factor=squeeze_factor,
                    conv_scale=conv_scale,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop,
                    attn_drop=attn_drop,
                    drop_path=drop_path[i]
                    if isinstance(drop_path, list)
                    else drop_path,
                    norm_layer=norm_layer,
                )
                for i in range(depth)
            ]
        )

        # OCAB
        self.overlap_attn = OCAB(
            dim=dim,
            input_resolution=input_resolution,
            window_size=window_size,
            overlap_ratio=overlap_ratio,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            mlp_ratio=mlp_ratio,  # type: ignore
            norm_layer=norm_layer,
        )

        # patch merging layer
        if downsample is not None:
            self.downsample = downsample(
                input_resolution, dim=dim, norm_layer=norm_layer
            )
        else:
            self.downsample = None

    def forward(self, x, x_size, params):
        for blk in self.blocks:
            x = blk(x, x_size, params["rpi_sa"], params["attn_mask"])

        x = self.overlap_attn(x, x_size, params["rpi_oca"])

        if self.downsample is not None:
            x = self.downsample(x)
        return x


class RHAG(nn.Module):
    """Residual Hybrid Attention Group (RHAG).
    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
        img_size: Input image size.
        patch_size: Patch size.
        resi_connection: The convolutional block before residual connection.
    """

    def __init__(
        self,
        dim,
        input_resolution,
        depth,
        num_heads,
        window_size,
        compress_ratio,
        squeeze_factor,
        conv_scale,
        overlap_ratio,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        downsample=None,
        use_checkpoint=False,
        img_size=224,
        patch_size=4,
        resi_connection="1conv",
    ):
        super(RHAG, self).__init__()

        self.dim = dim
        self.input_resolution = input_resolution

        self.residual_group = AttenBlocks(
            dim=dim,
            input_resolution=input_resolution,
            depth=depth,
            num_heads=num_heads,
            window_size=window_size,
            compress_ratio=compress_ratio,
            squeeze_factor=squeeze_factor,
            conv_scale=conv_scale,
            overlap_ratio=overlap_ratio,
            mlp_ratio=mlp_ratio,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            drop=drop,
            attn_drop=attn_drop,
            drop_path=drop_path,
            norm_layer=norm_layer,
            downsample=downsample,
            use_checkpoint=use_checkpoint,
        )

        if resi_connection == "1conv":
            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
        elif resi_connection == "identity":
            self.conv = nn.Identity()

        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=0,
            embed_dim=dim,
            norm_layer=None,
        )

        self.patch_unembed = PatchUnEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=0,
            embed_dim=dim,
            norm_layer=None,
        )

    def forward(self, x, x_size, params):
        return (
            self.patch_embed(
                self.conv(
                    self.patch_unembed(self.residual_group(x, x_size, params), x_size)
                )
            )
            + x
        )


class PatchEmbed(nn.Module):
    r"""Image to Patch Embedding
    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
        self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [
            img_size[0] // patch_size[0],  # type: ignore
            img_size[1] // patch_size[1],  # type: ignore
        ]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]

        self.in_chans = in_chans
        self.embed_dim = embed_dim

        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        x = x.flatten(2).transpose(1, 2)  # b Ph*Pw c
        if self.norm is not None:
            x = self.norm(x)
        return x


class PatchUnEmbed(nn.Module):
    r"""Image to Patch Unembedding
    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
        self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [
            img_size[0] // patch_size[0],  # type: ignore
            img_size[1] // patch_size[1],  # type: ignore
        ]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]

        self.in_chans = in_chans
        self.embed_dim = embed_dim

    def forward(self, x, x_size):
        x = (
            x.transpose(1, 2)
            .contiguous()
            .view(x.shape[0], self.embed_dim, x_size[0], x_size[1])
        )  # b Ph*Pw c
        return x


class Upsample(nn.Sequential):
    """Upsample module.
    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.
    """

    def __init__(self, scale, num_feat):
        m = []
        if (scale & (scale - 1)) == 0:  # scale = 2^n
            for _ in range(int(math.log(scale, 2))):
                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
                m.append(nn.PixelShuffle(2))
        elif scale == 3:
            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
            m.append(nn.PixelShuffle(3))
        else:
            raise ValueError(
                f"scale {scale} is not supported. " "Supported scales: 2^n and 3."
            )
        super(Upsample, self).__init__(*m)


class HAT(nn.Module):
    r"""Hybrid Attention Transformer
        A PyTorch implementation of : `Activating More Pixels in Image Super-Resolution Transformer`.
        Some codes are based on SwinIR.
    Args:
        img_size (int | tuple(int)): Input image size. Default 64
        patch_size (int | tuple(int)): Patch size. Default: 1
        in_chans (int): Number of input image channels. Default: 3
        embed_dim (int): Patch embedding dimension. Default: 96
        depths (tuple(int)): Depth of each Swin Transformer layer.
        num_heads (tuple(int)): Number of attention heads in different layers.
        window_size (int): Window size. Default: 7
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
        drop_rate (float): Dropout rate. Default: 0
        attn_drop_rate (float): Attention dropout rate. Default: 0
        drop_path_rate (float): Stochastic depth rate. Default: 0.1
        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
        patch_norm (bool): If True, add normalization after patch embedding. Default: True
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
        upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
        img_range: Image range. 1. or 255.
        upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
    """

    def __init__(
        self,
        state_dict,
        **kwargs,
    ):
        super(HAT, self).__init__()

        # Defaults
        img_size = 64
        patch_size = 1
        in_chans = 3
        embed_dim = 96
        depths = (6, 6, 6, 6)
        num_heads = (6, 6, 6, 6)
        window_size = 7
        compress_ratio = 3
        squeeze_factor = 30
        conv_scale = 0.01
        overlap_ratio = 0.5
        mlp_ratio = 4.0
        qkv_bias = True
        qk_scale = None
        drop_rate = 0.0
        attn_drop_rate = 0.0
        drop_path_rate = 0.1
        norm_layer = nn.LayerNorm
        ape = False
        patch_norm = True
        use_checkpoint = False
        upscale = 2
        img_range = 1.0
        upsampler = ""
        resi_connection = "1conv"

        self.state = state_dict
        self.model_arch = "HAT"
        self.sub_type = "SR"
        self.supports_fp16 = False
        self.support_bf16 = True
        self.min_size_restriction = 16

        state_keys = list(state_dict.keys())

        num_feat = state_dict["conv_last.weight"].shape[1]
        in_chans = state_dict["conv_first.weight"].shape[1]
        num_out_ch = state_dict["conv_last.weight"].shape[0]
        embed_dim = state_dict["conv_first.weight"].shape[0]

        if "conv_before_upsample.0.weight" in state_keys:
            if "conv_up1.weight" in state_keys:
                upsampler = "nearest+conv"
            else:
                upsampler = "pixelshuffle"
        elif "upsample.0.weight" in state_keys:
            upsampler = "pixelshuffledirect"
        else:
            upsampler = ""
        upscale = 1
        if upsampler == "nearest+conv":
            upsample_keys = [
                x for x in state_keys if "conv_up" in x and "bias" not in x
            ]

            for upsample_key in upsample_keys:
                upscale *= 2
        elif upsampler == "pixelshuffle":
            upsample_keys = [
                x
                for x in state_keys
                if "upsample" in x and "conv" not in x and "bias" not in x
            ]
            for upsample_key in upsample_keys:
                shape = self.state[upsample_key].shape[0]
                upscale *= math.sqrt(shape // num_feat)
            upscale = int(upscale)
        elif upsampler == "pixelshuffledirect":
            upscale = int(
                math.sqrt(self.state["upsample.0.bias"].shape[0] // num_out_ch)
            )

        max_layer_num = 0
        max_block_num = 0
        for key in state_keys:
            result = re.match(
                r"layers.(\d*).residual_group.blocks.(\d*).conv_block.cab.0.weight", key
            )
            if result:
                layer_num, block_num = result.groups()
                max_layer_num = max(max_layer_num, int(layer_num))
                max_block_num = max(max_block_num, int(block_num))

        depths = [max_block_num + 1 for _ in range(max_layer_num + 1)]

        if (
            "layers.0.residual_group.blocks.0.attn.relative_position_bias_table"
            in state_keys
        ):
            num_heads_num = self.state[
                "layers.0.residual_group.blocks.0.attn.relative_position_bias_table"
            ].shape[-1]
            num_heads = [num_heads_num for _ in range(max_layer_num + 1)]
        else:
            num_heads = depths

        mlp_ratio = float(
            self.state["layers.0.residual_group.blocks.0.mlp.fc1.bias"].shape[0]
            / embed_dim
        )

        # TODO: could actually count the layers, but this should do
        if "layers.0.conv.4.weight" in state_keys:
            resi_connection = "3conv"
        else:
            resi_connection = "1conv"

        window_size = int(math.sqrt(self.state["relative_position_index_SA"].shape[0]))

        # Not sure if this is needed or used at all anywhere in HAT's config
        if "layers.0.residual_group.blocks.1.attn_mask" in state_keys:
            img_size = int(
                math.sqrt(
                    self.state["layers.0.residual_group.blocks.1.attn_mask"].shape[0]
                )
                * window_size
            )

        self.window_size = window_size
        self.shift_size = window_size // 2
        self.overlap_ratio = overlap_ratio

        self.in_nc = in_chans
        self.out_nc = num_out_ch
        self.num_feat = num_feat
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.depths = depths
        self.window_size = window_size
        self.mlp_ratio = mlp_ratio
        self.scale = upscale
        self.upsampler = upsampler
        self.img_size = img_size
        self.img_range = img_range
        self.resi_connection = resi_connection

        num_in_ch = in_chans
        # num_out_ch = in_chans
        # num_feat = 64
        self.img_range = img_range
        if in_chans == 3:
            rgb_mean = (0.4488, 0.4371, 0.4040)
            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
        else:
            self.mean = torch.zeros(1, 1, 1, 1)
        self.upscale = upscale
        self.upsampler = upsampler

        # relative position index
        relative_position_index_SA = self.calculate_rpi_sa()
        relative_position_index_OCA = self.calculate_rpi_oca()
        self.register_buffer("relative_position_index_SA", relative_position_index_SA)
        self.register_buffer("relative_position_index_OCA", relative_position_index_OCA)

        # ------------------------- 1, shallow feature extraction ------------------------- #
        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)

        # ------------------------- 2, deep feature extraction ------------------------- #
        self.num_layers = len(depths)
        self.embed_dim = embed_dim
        self.ape = ape
        self.patch_norm = patch_norm
        self.num_features = embed_dim
        self.mlp_ratio = mlp_ratio

        # split image into non-overlapping patches
        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=embed_dim,
            embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None,
        )
        num_patches = self.patch_embed.num_patches
        patches_resolution = self.patch_embed.patches_resolution
        self.patches_resolution = patches_resolution

        # merge non-overlapping patches into image
        self.patch_unembed = PatchUnEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=embed_dim,
            embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None,
        )

        # absolute position embedding
        if self.ape:
            self.absolute_pos_embed = nn.Parameter(  # type: ignore[arg-type]
                torch.zeros(1, num_patches, embed_dim)
            )
            trunc_normal_(self.absolute_pos_embed, std=0.02)

        self.pos_drop = nn.Dropout(p=drop_rate)

        # stochastic depth
        dpr = [
            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
        ]  # stochastic depth decay rule

        # build Residual Hybrid Attention Groups (RHAG)
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):
            layer = RHAG(
                dim=embed_dim,
                input_resolution=(patches_resolution[0], patches_resolution[1]),
                depth=depths[i_layer],
                num_heads=num_heads[i_layer],
                window_size=window_size,
                compress_ratio=compress_ratio,
                squeeze_factor=squeeze_factor,
                conv_scale=conv_scale,
                overlap_ratio=overlap_ratio,
                mlp_ratio=self.mlp_ratio,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop=drop_rate,
                attn_drop=attn_drop_rate,
                drop_path=dpr[
                    sum(depths[:i_layer]) : sum(depths[: i_layer + 1])  # type: ignore
                ],  # no impact on SR results
                norm_layer=norm_layer,
                downsample=None,
                use_checkpoint=use_checkpoint,
                img_size=img_size,
                patch_size=patch_size,
                resi_connection=resi_connection,
            )
            self.layers.append(layer)
        self.norm = norm_layer(self.num_features)

        # build the last conv layer in deep feature extraction
        if resi_connection == "1conv":
            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
        elif resi_connection == "identity":
            self.conv_after_body = nn.Identity()

        # ------------------------- 3, high quality image reconstruction ------------------------- #
        if self.upsampler == "pixelshuffle":
            # for classical SR
            self.conv_before_upsample = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.upsample = Upsample(upscale, num_feat)
            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

        self.apply(self._init_weights)
        self.load_state_dict(self.state, strict=False)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def calculate_rpi_sa(self):
        # calculate relative position index for SA
        coords_h = torch.arange(self.window_size)
        coords_w = torch.arange(self.window_size)
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
        relative_coords = (
            coords_flatten[:, :, None] - coords_flatten[:, None, :]
        )  # 2, Wh*Ww, Wh*Ww
        relative_coords = relative_coords.permute(
            1, 2, 0
        ).contiguous()  # Wh*Ww, Wh*Ww, 2
        relative_coords[:, :, 0] += self.window_size - 1  # shift to start from 0
        relative_coords[:, :, 1] += self.window_size - 1
        relative_coords[:, :, 0] *= 2 * self.window_size - 1
        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
        return relative_position_index

    def calculate_rpi_oca(self):
        # calculate relative position index for OCA
        window_size_ori = self.window_size
        window_size_ext = self.window_size + int(self.overlap_ratio * self.window_size)

        coords_h = torch.arange(window_size_ori)
        coords_w = torch.arange(window_size_ori)
        coords_ori = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, ws, ws
        coords_ori_flatten = torch.flatten(coords_ori, 1)  # 2, ws*ws

        coords_h = torch.arange(window_size_ext)
        coords_w = torch.arange(window_size_ext)
        coords_ext = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, wse, wse
        coords_ext_flatten = torch.flatten(coords_ext, 1)  # 2, wse*wse

        relative_coords = (
            coords_ext_flatten[:, None, :] - coords_ori_flatten[:, :, None]
        )  # 2, ws*ws, wse*wse

        relative_coords = relative_coords.permute(
            1, 2, 0
        ).contiguous()  # ws*ws, wse*wse, 2
        relative_coords[:, :, 0] += (
            window_size_ori - window_size_ext + 1
        )  # shift to start from 0
        relative_coords[:, :, 1] += window_size_ori - window_size_ext + 1

        relative_coords[:, :, 0] *= window_size_ori + window_size_ext - 1
        relative_position_index = relative_coords.sum(-1)
        return relative_position_index

    def calculate_mask(self, x_size):
        # calculate attention mask for SW-MSA
        h, w = x_size
        img_mask = torch.zeros((1, h, w, 1))  # 1 h w 1
        h_slices = (
            slice(0, -self.window_size),
            slice(-self.window_size, -self.shift_size),
            slice(-self.shift_size, None),
        )
        w_slices = (
            slice(0, -self.window_size),
            slice(-self.window_size, -self.shift_size),
            slice(-self.shift_size, None),
        )
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                img_mask[:, h, w, :] = cnt
                cnt += 1

        mask_windows = window_partition(
            img_mask, self.window_size
        )  # nw, window_size, window_size, 1
        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
            attn_mask == 0, float(0.0)
        )

        return attn_mask

    @torch.jit.ignore  # type: ignore
    def no_weight_decay(self):
        return {"absolute_pos_embed"}

    @torch.jit.ignore  # type: ignore
    def no_weight_decay_keywords(self):
        return {"relative_position_bias_table"}

    def check_image_size(self, x):
        _, _, h, w = x.size()
        mod_pad_h = (self.window_size - h % self.window_size) % self.window_size
        mod_pad_w = (self.window_size - w % self.window_size) % self.window_size
        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect")
        return x

    def forward_features(self, x):
        x_size = (x.shape[2], x.shape[3])

        # Calculate attention mask and relative position index in advance to speed up inference.
        # The original code is very time-cosuming for large window size.
        attn_mask = self.calculate_mask(x_size).to(x.device)
        params = {
            "attn_mask": attn_mask,
            "rpi_sa": self.relative_position_index_SA,
            "rpi_oca": self.relative_position_index_OCA,
        }

        x = self.patch_embed(x)
        if self.ape:
            x = x + self.absolute_pos_embed
        x = self.pos_drop(x)

        for layer in self.layers:
            x = layer(x, x_size, params)

        x = self.norm(x)  # b seq_len c
        x = self.patch_unembed(x, x_size)

        return x

    def forward(self, x):
        H, W = x.shape[2:]
        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range
        x = self.check_image_size(x)

        if self.upsampler == "pixelshuffle":
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))

        x = x / self.img_range + self.mean

        return x[:, :, : H * self.upscale, : W * self.upscale]

nodes/impl/pytorch/architecture/HAT.py
vladmandic/sd-extension-chainner
# pylint: skip-file
"""
Model adapted from advimman's lama project: https://github.com/advimman/lama
"""

# Fast Fourier Convolution NeurIPS 2020
# original implementation https://github.com/pkumivision/FFC/blob/main/model_zoo/ffc.py
# paper https://proceedings.neurips.cc/paper/2020/file/2fd5d41ec6cfab47e32164d5624269b1-Paper.pdf

from typing import List

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.transforms.functional import InterpolationMode, rotate


class LearnableSpatialTransformWrapper(nn.Module):
    def __init__(self, impl, pad_coef=0.5, angle_init_range=80, train_angle=True):
        super().__init__()
        self.impl = impl
        self.angle = torch.rand(1) * angle_init_range
        if train_angle:
            self.angle = nn.Parameter(self.angle, requires_grad=True)
        self.pad_coef = pad_coef

    def forward(self, x):
        if torch.is_tensor(x):
            return self.inverse_transform(self.impl(self.transform(x)), x)
        elif isinstance(x, tuple):
            x_trans = tuple(self.transform(elem) for elem in x)
            y_trans = self.impl(x_trans)
            return tuple(
                self.inverse_transform(elem, orig_x) for elem, orig_x in zip(y_trans, x)
            )
        else:
            raise ValueError(f"Unexpected input type {type(x)}")

    def transform(self, x):
        height, width = x.shape[2:]
        pad_h, pad_w = int(height * self.pad_coef), int(width * self.pad_coef)
        x_padded = F.pad(x, [pad_w, pad_w, pad_h, pad_h], mode="reflect")
        x_padded_rotated = rotate(
            x_padded, self.angle.to(x_padded), InterpolationMode.BILINEAR, fill=0
        )

        return x_padded_rotated

    def inverse_transform(self, y_padded_rotated, orig_x):
        height, width = orig_x.shape[2:]
        pad_h, pad_w = int(height * self.pad_coef), int(width * self.pad_coef)

        y_padded = rotate(
            y_padded_rotated,
            -self.angle.to(y_padded_rotated),
            InterpolationMode.BILINEAR,
            fill=0,
        )
        y_height, y_width = y_padded.shape[2:]
        y = y_padded[:, :, pad_h : y_height - pad_h, pad_w : y_width - pad_w]
        return y


class SELayer(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid(),
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        res = x * y.expand_as(x)
        return res


class FourierUnit(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        groups=1,
        spatial_scale_factor=None,
        spatial_scale_mode="bilinear",
        spectral_pos_encoding=False,
        use_se=False,
        se_kwargs=None,
        ffc3d=False,
        fft_norm="ortho",
    ):
        # bn_layer not used
        super(FourierUnit, self).__init__()
        self.groups = groups

        self.conv_layer = torch.nn.Conv2d(
            in_channels=in_channels * 2 + (2 if spectral_pos_encoding else 0),
            out_channels=out_channels * 2,
            kernel_size=1,
            stride=1,
            padding=0,
            groups=self.groups,
            bias=False,
        )
        self.bn = torch.nn.BatchNorm2d(out_channels * 2)
        self.relu = torch.nn.ReLU(inplace=True)

        # squeeze and excitation block
        self.use_se = use_se
        if use_se:
            if se_kwargs is None:
                se_kwargs = {}
            self.se = SELayer(self.conv_layer.in_channels, **se_kwargs)

        self.spatial_scale_factor = spatial_scale_factor
        self.spatial_scale_mode = spatial_scale_mode
        self.spectral_pos_encoding = spectral_pos_encoding
        self.ffc3d = ffc3d
        self.fft_norm = fft_norm

    def forward(self, x):
        half_check = False
        if x.type() == "torch.cuda.HalfTensor":
            # half only works on gpu anyway
            half_check = True

        batch = x.shape[0]

        if self.spatial_scale_factor is not None:
            orig_size = x.shape[-2:]
            x = F.interpolate(
                x,
                scale_factor=self.spatial_scale_factor,
                mode=self.spatial_scale_mode,
                align_corners=False,
            )

        # (batch, c, h, w/2+1, 2)
        fft_dim = (-3, -2, -1) if self.ffc3d else (-2, -1)
        if half_check is True:
            ffted = torch.fft.rfftn(
                x.float(), dim=fft_dim, norm=self.fft_norm
            )  # .type(torch.cuda.HalfTensor)
        else:
            ffted = torch.fft.rfftn(x, dim=fft_dim, norm=self.fft_norm)

        ffted = torch.stack((ffted.real, ffted.imag), dim=-1)
        ffted = ffted.permute(0, 1, 4, 2, 3).contiguous()  # (batch, c, 2, h, w/2+1)
        ffted = ffted.view(
            (
                batch,
                -1,
            )
            + ffted.size()[3:]
        )

        if self.spectral_pos_encoding:
            height, width = ffted.shape[-2:]
            coords_vert = (
                torch.linspace(0, 1, height)[None, None, :, None]
                .expand(batch, 1, height, width)
                .to(ffted)
            )
            coords_hor = (
                torch.linspace(0, 1, width)[None, None, None, :]
                .expand(batch, 1, height, width)
                .to(ffted)
            )
            ffted = torch.cat((coords_vert, coords_hor, ffted), dim=1)

        if self.use_se:
            ffted = self.se(ffted)

        if half_check is True:
            ffted = self.conv_layer(ffted.half())  # (batch, c*2, h, w/2+1)
        else:
            ffted = self.conv_layer(
                ffted
            )  # .type(torch.cuda.FloatTensor)  # (batch, c*2, h, w/2+1)

        ffted = self.relu(self.bn(ffted))
        # forcing to be always float
        ffted = ffted.float()

        ffted = (
            ffted.view(
                (
                    batch,
                    -1,
                    2,
                )
                + ffted.size()[2:]
            )
            .permute(0, 1, 3, 4, 2)
            .contiguous()
        )  # (batch,c, t, h, w/2+1, 2)

        ffted = torch.complex(ffted[..., 0], ffted[..., 1])

        ifft_shape_slice = x.shape[-3:] if self.ffc3d else x.shape[-2:]
        output = torch.fft.irfftn(
            ffted, s=ifft_shape_slice, dim=fft_dim, norm=self.fft_norm
        )

        if half_check is True:
            output = output.half()

        if self.spatial_scale_factor is not None:
            output = F.interpolate(
                output,
                size=orig_size,
                mode=self.spatial_scale_mode,
                align_corners=False,
            )

        return output


class SpectralTransform(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        stride=1,
        groups=1,
        enable_lfu=True,
        separable_fu=False,
        **fu_kwargs,
    ):
        # bn_layer not used
        super(SpectralTransform, self).__init__()
        self.enable_lfu = enable_lfu
        if stride == 2:
            self.downsample = nn.AvgPool2d(kernel_size=(2, 2), stride=2)
        else:
            self.downsample = nn.Identity()

        self.stride = stride
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels, out_channels // 2, kernel_size=1, groups=groups, bias=False
            ),
            nn.BatchNorm2d(out_channels // 2),
            nn.ReLU(inplace=True),
        )
        fu_class = FourierUnit
        self.fu = fu_class(out_channels // 2, out_channels // 2, groups, **fu_kwargs)
        if self.enable_lfu:
            self.lfu = fu_class(out_channels // 2, out_channels // 2, groups)
        self.conv2 = torch.nn.Conv2d(
            out_channels // 2, out_channels, kernel_size=1, groups=groups, bias=False
        )

    def forward(self, x):
        x = self.downsample(x)
        x = self.conv1(x)
        output = self.fu(x)

        if self.enable_lfu:
            _, c, h, _ = x.shape
            split_no = 2
            split_s = h // split_no
            xs = torch.cat(
                torch.split(x[:, : c // 4], split_s, dim=-2), dim=1
            ).contiguous()
            xs = torch.cat(torch.split(xs, split_s, dim=-1), dim=1).contiguous()
            xs = self.lfu(xs)
            xs = xs.repeat(1, 1, split_no, split_no).contiguous()
        else:
            xs = 0

        output = self.conv2(x + output + xs)

        return output


class FFC(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        ratio_gin,
        ratio_gout,
        stride=1,
        padding=0,
        dilation=1,
        groups=1,
        bias=False,
        enable_lfu=True,
        padding_type="reflect",
        gated=False,
        **spectral_kwargs,
    ):
        super(FFC, self).__init__()

        assert stride == 1 or stride == 2, "Stride should be 1 or 2."
        self.stride = stride

        in_cg = int(in_channels * ratio_gin)
        in_cl = in_channels - in_cg
        out_cg = int(out_channels * ratio_gout)
        out_cl = out_channels - out_cg
        # groups_g = 1 if groups == 1 else int(groups * ratio_gout)
        # groups_l = 1 if groups == 1 else groups - groups_g

        self.ratio_gin = ratio_gin
        self.ratio_gout = ratio_gout
        self.global_in_num = in_cg

        module = nn.Identity if in_cl == 0 or out_cl == 0 else nn.Conv2d
        self.convl2l = module(
            in_cl,
            out_cl,
            kernel_size,
            stride,
            padding,
            dilation,
            groups,
            bias,
            padding_mode=padding_type,
        )
        module = nn.Identity if in_cl == 0 or out_cg == 0 else nn.Conv2d
        self.convl2g = module(
            in_cl,
            out_cg,
            kernel_size,
            stride,
            padding,
            dilation,
            groups,
            bias,
            padding_mode=padding_type,
        )
        module = nn.Identity if in_cg == 0 or out_cl == 0 else nn.Conv2d
        self.convg2l = module(
            in_cg,
            out_cl,
            kernel_size,
            stride,
            padding,
            dilation,
            groups,
            bias,
            padding_mode=padding_type,
        )
        module = nn.Identity if in_cg == 0 or out_cg == 0 else SpectralTransform
        self.convg2g = module(
            in_cg,
            out_cg,
            stride,
            1 if groups == 1 else groups // 2,
            enable_lfu,
            **spectral_kwargs,
        )

        self.gated = gated
        module = (
            nn.Identity if in_cg == 0 or out_cl == 0 or not self.gated else nn.Conv2d
        )
        self.gate = module(in_channels, 2, 1)

    def forward(self, x):
        x_l, x_g = x if type(x) is tuple else (x, 0)
        out_xl, out_xg = 0, 0

        if self.gated:
            total_input_parts = [x_l]
            if torch.is_tensor(x_g):
                total_input_parts.append(x_g)
            total_input = torch.cat(total_input_parts, dim=1)

            gates = torch.sigmoid(self.gate(total_input))
            g2l_gate, l2g_gate = gates.chunk(2, dim=1)
        else:
            g2l_gate, l2g_gate = 1, 1

        if self.ratio_gout != 1:
            out_xl = self.convl2l(x_l) + self.convg2l(x_g) * g2l_gate
        if self.ratio_gout != 0:
            out_xg = self.convl2g(x_l) * l2g_gate + self.convg2g(x_g)

        return out_xl, out_xg


class FFC_BN_ACT(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        ratio_gin,
        ratio_gout,
        stride=1,
        padding=0,
        dilation=1,
        groups=1,
        bias=False,
        norm_layer=nn.BatchNorm2d,
        activation_layer=nn.Identity,
        padding_type="reflect",
        enable_lfu=True,
        **kwargs,
    ):
        super(FFC_BN_ACT, self).__init__()
        self.ffc = FFC(
            in_channels,
            out_channels,
            kernel_size,
            ratio_gin,
            ratio_gout,
            stride,
            padding,
            dilation,
            groups,
            bias,
            enable_lfu,
            padding_type=padding_type,
            **kwargs,
        )
        lnorm = nn.Identity if ratio_gout == 1 else norm_layer
        gnorm = nn.Identity if ratio_gout == 0 else norm_layer
        global_channels = int(out_channels * ratio_gout)
        self.bn_l = lnorm(out_channels - global_channels)
        self.bn_g = gnorm(global_channels)

        lact = nn.Identity if ratio_gout == 1 else activation_layer
        gact = nn.Identity if ratio_gout == 0 else activation_layer
        self.act_l = lact(inplace=True)
        self.act_g = gact(inplace=True)

    def forward(self, x):
        x_l, x_g = self.ffc(x)
        x_l = self.act_l(self.bn_l(x_l))
        x_g = self.act_g(self.bn_g(x_g))
        return x_l, x_g


class FFCResnetBlock(nn.Module):
    def __init__(
        self,
        dim,
        padding_type,
        norm_layer,
        activation_layer=nn.ReLU,
        dilation=1,
        spatial_transform_kwargs=None,
        inline=False,
        **conv_kwargs,
    ):
        super().__init__()
        self.conv1 = FFC_BN_ACT(
            dim,
            dim,
            kernel_size=3,
            padding=dilation,
            dilation=dilation,
            norm_layer=norm_layer,
            activation_layer=activation_layer,
            padding_type=padding_type,
            **conv_kwargs,
        )
        self.conv2 = FFC_BN_ACT(
            dim,
            dim,
            kernel_size=3,
            padding=dilation,
            dilation=dilation,
            norm_layer=norm_layer,
            activation_layer=activation_layer,
            padding_type=padding_type,
            **conv_kwargs,
        )
        if spatial_transform_kwargs is not None:
            self.conv1 = LearnableSpatialTransformWrapper(
                self.conv1, **spatial_transform_kwargs
            )
            self.conv2 = LearnableSpatialTransformWrapper(
                self.conv2, **spatial_transform_kwargs
            )
        self.inline = inline

    def forward(self, x):
        if self.inline:
            x_l, x_g = (
                x[:, : -self.conv1.ffc.global_in_num],
                x[:, -self.conv1.ffc.global_in_num :],
            )
        else:
            x_l, x_g = x if type(x) is tuple else (x, 0)

        id_l, id_g = x_l, x_g

        x_l, x_g = self.conv1((x_l, x_g))
        x_l, x_g = self.conv2((x_l, x_g))

        x_l, x_g = id_l + x_l, id_g + x_g
        out = x_l, x_g
        if self.inline:
            out = torch.cat(out, dim=1)
        return out


class ConcatTupleLayer(nn.Module):
    def forward(self, x):
        assert isinstance(x, tuple)
        x_l, x_g = x
        assert torch.is_tensor(x_l) or torch.is_tensor(x_g)
        if not torch.is_tensor(x_g):
            return x_l
        return torch.cat(x, dim=1)


class FFCResNetGenerator(nn.Module):
    def __init__(
        self,
        input_nc,
        output_nc,
        ngf=64,
        n_downsampling=3,
        n_blocks=18,
        norm_layer=nn.BatchNorm2d,
        padding_type="reflect",
        activation_layer=nn.ReLU,
        up_norm_layer=nn.BatchNorm2d,
        up_activation=nn.ReLU(True),
        init_conv_kwargs=None,
        downsample_conv_kwargs=None,
        resnet_conv_kwargs=None,
        spatial_transform_layers=None,
        spatial_transform_kwargs=None,
        max_features=1024,
        out_ffc=False,
        out_ffc_kwargs=None,
    ):
        if out_ffc_kwargs is None:
            out_ffc_kwargs = {}
        if spatial_transform_kwargs is None:
            spatial_transform_kwargs = {}
        if resnet_conv_kwargs is None:
            resnet_conv_kwargs = {}
        if downsample_conv_kwargs is None:
            downsample_conv_kwargs = {}
        if init_conv_kwargs is None:
            init_conv_kwargs = {}
        assert n_blocks >= 0
        super().__init__()
        """
        init_conv_kwargs = {'ratio_gin': 0, 'ratio_gout': 0, 'enable_lfu': False}
        downsample_conv_kwargs = {'ratio_gin': '${generator.init_conv_kwargs.ratio_gout}', 'ratio_gout': '${generator.downsample_conv_kwargs.ratio_gin}', 'enable_lfu': False}
        resnet_conv_kwargs = {'ratio_gin': 0.75, 'ratio_gout': '${generator.resnet_conv_kwargs.ratio_gin}', 'enable_lfu': False}
        spatial_transform_kwargs = {}
        out_ffc_kwargs = {}
        """
        """
        print(input_nc, output_nc, ngf, n_downsampling, n_blocks, norm_layer,
                padding_type, activation_layer,
                up_norm_layer, up_activation,
                spatial_transform_layers,
                add_out_act, max_features, out_ffc, file=sys.stderr)

        4 3 64 3 18 <class 'torch.nn.modules.batchnorm.BatchNorm2d'>
        reflect <class 'torch.nn.modules.activation.ReLU'>
        <class 'torch.nn.modules.batchnorm.BatchNorm2d'>
        ReLU(inplace=True)
        None sigmoid 1024 False
        """
        init_conv_kwargs = {"ratio_gin": 0, "ratio_gout": 0, "enable_lfu": False}
        downsample_conv_kwargs = {"ratio_gin": 0, "ratio_gout": 0, "enable_lfu": False}
        resnet_conv_kwargs = {
            "ratio_gin": 0.75,
            "ratio_gout": 0.75,
            "enable_lfu": False,
        }
        spatial_transform_kwargs = {}
        out_ffc_kwargs = {}

        model = [
            nn.ReflectionPad2d(3),
            FFC_BN_ACT(
                input_nc,
                ngf,
                kernel_size=7,
                padding=0,
                norm_layer=norm_layer,
                activation_layer=activation_layer,
                **init_conv_kwargs,
            ),
        ]

        ### downsample
        for i in range(n_downsampling):
            mult = 2**i
            if i == n_downsampling - 1:
                cur_conv_kwargs = dict(downsample_conv_kwargs)
                cur_conv_kwargs["ratio_gout"] = resnet_conv_kwargs.get("ratio_gin", 0)
            else:
                cur_conv_kwargs = downsample_conv_kwargs
            model += [
                FFC_BN_ACT(
                    min(max_features, ngf * mult),
                    min(max_features, ngf * mult * 2),
                    kernel_size=3,
                    stride=2,
                    padding=1,
                    norm_layer=norm_layer,
                    activation_layer=activation_layer,
                    **cur_conv_kwargs,
                )
            ]

        mult = 2**n_downsampling
        feats_num_bottleneck = min(max_features, ngf * mult)

        ### resnet blocks
        for i in range(n_blocks):
            cur_resblock = FFCResnetBlock(
                feats_num_bottleneck,
                padding_type=padding_type,
                activation_layer=activation_layer,
                norm_layer=norm_layer,
                **resnet_conv_kwargs,
            )
            if spatial_transform_layers is not None and i in spatial_transform_layers:
                cur_resblock = LearnableSpatialTransformWrapper(
                    cur_resblock, **spatial_transform_kwargs
                )
            model += [cur_resblock]

        model += [ConcatTupleLayer()]

        ### upsample
        for i in range(n_downsampling):
            mult = 2 ** (n_downsampling - i)
            model += [
                nn.ConvTranspose2d(
                    min(max_features, ngf * mult),
                    min(max_features, int(ngf * mult / 2)),
                    kernel_size=3,
                    stride=2,
                    padding=1,
                    output_padding=1,
                ),
                up_norm_layer(min(max_features, int(ngf * mult / 2))),
                up_activation,
            ]

        if out_ffc:
            model += [
                FFCResnetBlock(
                    ngf,
                    padding_type=padding_type,
                    activation_layer=activation_layer,
                    norm_layer=norm_layer,
                    inline=True,
                    **out_ffc_kwargs,
                )
            ]

        model += [
            nn.ReflectionPad2d(3),
            nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
        ]
        model.append(nn.Sigmoid())
        self.model = nn.Sequential(*model)

    def forward(self, image, mask):
        return self.model(torch.cat([image, mask], dim=1))


class LaMa(nn.Module):
    def __init__(self, state_dict) -> None:
        super(LaMa, self).__init__()
        self.model_arch = "LaMa"
        self.sub_type = "Inpaint"
        self.in_nc = 4
        self.out_nc = 3
        self.scale = 1

        self.min_size = None
        self.pad_mod = 8
        self.pad_to_square = False

        self.model = FFCResNetGenerator(self.in_nc, self.out_nc)
        self.state = {
            k.replace("generator.model", "model.model"): v
            for k, v in state_dict.items()
        }

        self.supports_fp16 = False
        self.support_bf16 = True

        self.load_state_dict(self.state, strict=False)

    def forward(self, img, mask):
        masked_img = img * (1 - mask)
        inpainted_mask = mask * self.model.forward(masked_img, mask)
        result = inpainted_mask + (1 - mask) * img
        return result

nodes/impl/pytorch/architecture/LaMa.py
vladmandic/sd-extension-chainner
# pylint: skip-file
"""Original MAT project is copyright of fenglingwb: https://github.com/fenglinglwb/MAT
Code used for this implementation of MAT is modified from lama-cleaner,
copyright of Sanster: https://github.com/fenglinglwb/MAT"""

import random

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint

from .mat.utils import (
    Conv2dLayer,
    FullyConnectedLayer,
    activation_funcs,
    bias_act,
    conv2d_resample,
    normalize_2nd_moment,
    setup_filter,
    to_2tuple,
    upsample2d,
)


class ModulatedConv2d(nn.Module):
    def __init__(
        self,
        in_channels,  # Number of input channels.
        out_channels,  # Number of output channels.
        kernel_size,  # Width and height of the convolution kernel.
        style_dim,  # dimension of the style code
        demodulate=True,  # perfrom demodulation
        up=1,  # Integer upsampling factor.
        down=1,  # Integer downsampling factor.
        resample_filter=None,  # Low-pass filter to apply when resampling activations.
        conv_clamp=None,  # Clamp the output to +-X, None = disable clamping.
    ):
        if resample_filter is None:
            resample_filter = [1, 3, 3, 1]
        super().__init__()
        self.demodulate = demodulate

        self.weight = torch.nn.Parameter(
            torch.randn([1, out_channels, in_channels, kernel_size, kernel_size])
        )
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size**2))
        self.padding = self.kernel_size // 2
        self.up = up
        self.down = down
        self.register_buffer("resample_filter", setup_filter(resample_filter))
        self.conv_clamp = conv_clamp

        self.affine = FullyConnectedLayer(style_dim, in_channels, bias_init=1)

    def forward(self, x, style):
        batch, in_channels, height, width = x.shape
        style = self.affine(style).view(batch, 1, in_channels, 1, 1).to(x.device)
        weight = self.weight.to(x.device) * self.weight_gain * style

        if self.demodulate:
            decoefs = (weight.pow(2).sum(dim=[2, 3, 4]) + 1e-8).rsqrt()
            weight = weight * decoefs.view(batch, self.out_channels, 1, 1, 1)

        weight = weight.view(
            batch * self.out_channels, in_channels, self.kernel_size, self.kernel_size
        )
        x = x.view(1, batch * in_channels, height, width)
        x = conv2d_resample(
            x=x,
            w=weight,
            f=self.resample_filter,
            up=self.up,
            down=self.down,
            padding=self.padding,
            groups=batch,
        )
        out = x.view(batch, self.out_channels, *x.shape[2:])

        return out


class StyleConv(torch.nn.Module):
    def __init__(
        self,
        in_channels,  # Number of input channels.
        out_channels,  # Number of output channels.
        style_dim,  # Intermediate latent (W) dimensionality.
        resolution,  # Resolution of this layer.
        kernel_size=3,  # Convolution kernel size.
        up=1,  # Integer upsampling factor.
        use_noise=False,  # Enable noise input?
        activation="lrelu",  # Activation function: 'relu', 'lrelu', etc.
        resample_filter=None,  # Low-pass filter to apply when resampling activations.
        conv_clamp=None,  # Clamp the output of convolution layers to +-X, None = disable clamping.
        demodulate=True,  # perform demodulation
    ):
        if resample_filter is None:
            resample_filter = [1, 3, 3, 1]
        super().__init__()

        self.conv = ModulatedConv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            style_dim=style_dim,
            demodulate=demodulate,
            up=up,
            resample_filter=resample_filter,
            conv_clamp=conv_clamp,
        )

        self.use_noise = use_noise
        self.resolution = resolution
        if use_noise:
            self.register_buffer("noise_const", torch.randn([resolution, resolution]))
            self.noise_strength = torch.nn.Parameter(torch.zeros([]))

        self.bias = torch.nn.Parameter(torch.zeros([out_channels]))
        self.activation = activation
        self.act_gain = activation_funcs[activation].def_gain
        self.conv_clamp = conv_clamp

    def forward(self, x, style, noise_mode="random", gain=1):
        x = self.conv(x, style)

        assert noise_mode in ["random", "const", "none"]

        if self.use_noise:
            if noise_mode == "random":
                xh, xw = x.size()[-2:]
                noise = (
                    torch.randn([x.shape[0], 1, xh, xw], device=x.device)
                    * self.noise_strength
                )
            if noise_mode == "const":
                noise = self.noise_const * self.noise_strength
            x = x + noise

        act_gain = self.act_gain * gain
        act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None
        out = bias_act(
            x, self.bias, act=self.activation, gain=act_gain, clamp=act_clamp
        )

        return out


class ToRGB(torch.nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        style_dim,
        kernel_size=1,
        resample_filter=None,
        conv_clamp=None,
        demodulate=False,
    ):
        if resample_filter is None:
            resample_filter = [1, 3, 3, 1]
        super().__init__()

        self.conv = ModulatedConv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            style_dim=style_dim,
            demodulate=demodulate,
            resample_filter=resample_filter,
            conv_clamp=conv_clamp,
        )
        self.bias = torch.nn.Parameter(torch.zeros([out_channels]))
        self.register_buffer("resample_filter", setup_filter(resample_filter))
        self.conv_clamp = conv_clamp

    def forward(self, x, style, skip=None):
        x = self.conv(x, style)
        out = bias_act(x, self.bias, clamp=self.conv_clamp)

        if skip is not None:
            if skip.shape != out.shape:
                skip = upsample2d(skip, self.resample_filter)
            out = out + skip

        return out


def get_style_code(a, b):
    return torch.cat([a, b.to(a.device)], dim=1)


class DecBlockFirst(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        activation,
        style_dim,
        use_noise,
        demodulate,
        img_channels,
    ):
        super().__init__()
        self.fc = FullyConnectedLayer(
            in_features=in_channels * 2,
            out_features=in_channels * 4**2,
            activation=activation,
        )
        self.conv = StyleConv(
            in_channels=in_channels,
            out_channels=out_channels,
            style_dim=style_dim,
            resolution=4,
            kernel_size=3,
            use_noise=use_noise,
            activation=activation,
            demodulate=demodulate,
        )
        self.toRGB = ToRGB(
            in_channels=out_channels,
            out_channels=img_channels,
            style_dim=style_dim,
            kernel_size=1,
            demodulate=False,
        )

    def forward(self, x, ws, gs, E_features, noise_mode="random"):
        x = self.fc(x).view(x.shape[0], -1, 4, 4)
        x = x + E_features[2]
        style = get_style_code(ws[:, 0], gs)
        x = self.conv(x, style, noise_mode=noise_mode)
        style = get_style_code(ws[:, 1], gs)
        img = self.toRGB(x, style, skip=None)

        return x, img


class MappingNet(torch.nn.Module):
    def __init__(
        self,
        z_dim,  # Input latent (Z) dimensionality, 0 = no latent.
        c_dim,  # Conditioning label (C) dimensionality, 0 = no label.
        w_dim,  # Intermediate latent (W) dimensionality.
        num_ws,  # Number of intermediate latents to output, None = do not broadcast.
        num_layers=8,  # Number of mapping layers.
        embed_features=None,  # Label embedding dimensionality, None = same as w_dim.
        layer_features=None,  # Number of intermediate features in the mapping layers, None = same as w_dim.
        activation="lrelu",  # Activation function: 'relu', 'lrelu', etc.
        lr_multiplier=0.01,  # Learning rate multiplier for the mapping layers.
        w_avg_beta=0.995,  # Decay for tracking the moving average of W during training, None = do not track.
    ):
        super().__init__()
        self.z_dim = z_dim
        self.c_dim = c_dim
        self.w_dim = w_dim
        self.num_ws = num_ws
        self.num_layers = num_layers
        self.w_avg_beta = w_avg_beta

        if embed_features is None:
            embed_features = w_dim
        if c_dim == 0:
            embed_features = 0
        if layer_features is None:
            layer_features = w_dim
        features_list = (
            [z_dim + embed_features] + [layer_features] * (num_layers - 1) + [w_dim]
        )

        if c_dim > 0:
            self.embed = FullyConnectedLayer(c_dim, embed_features)
        for idx in range(num_layers):
            in_features = features_list[idx]
            out_features = features_list[idx + 1]
            layer = FullyConnectedLayer(
                in_features,
                out_features,
                activation=activation,
                lr_multiplier=lr_multiplier,
            )
            setattr(self, f"fc{idx}", layer)

        if num_ws is not None and w_avg_beta is not None:
            self.register_buffer("w_avg", torch.zeros([w_dim]))

    def forward(
        self, z, c, truncation_psi=1, truncation_cutoff=None, skip_w_avg_update=False
    ):
        # Embed, normalize, and concat inputs.
        x = None
        with torch.autograd.profiler.record_function("input"):
            if self.z_dim > 0:
                x = normalize_2nd_moment(z.to(torch.float32))
            if self.c_dim > 0:
                y = normalize_2nd_moment(self.embed(c.to(torch.float32)))
                x = torch.cat([x, y], dim=1) if x is not None else y

        # Main layers.
        for idx in range(self.num_layers):
            layer = getattr(self, f"fc{idx}")
            x = layer(x)

        # Update moving average of W.
        if self.w_avg_beta is not None and self.training and not skip_w_avg_update:
            with torch.autograd.profiler.record_function("update_w_avg"):
                self.w_avg.copy_(
                    x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta)
                )

        # Broadcast.
        if self.num_ws is not None:
            with torch.autograd.profiler.record_function("broadcast"):
                x = x.unsqueeze(1).repeat([1, self.num_ws, 1])

        # Apply truncation.
        if truncation_psi != 1:
            with torch.autograd.profiler.record_function("truncate"):
                assert self.w_avg_beta is not None
                if self.num_ws is None or truncation_cutoff is None:
                    x = self.w_avg.lerp(x, truncation_psi)
                else:
                    x[:, :truncation_cutoff] = self.w_avg.lerp(
                        x[:, :truncation_cutoff], truncation_psi
                    )

        return x


class DisFromRGB(nn.Module):
    def __init__(
        self, in_channels, out_channels, activation
    ):  # res = 2, ..., resolution_log2
        super().__init__()
        self.conv = Conv2dLayer(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=1,
            activation=activation,
        )

    def forward(self, x):
        return self.conv(x)


class DisBlock(nn.Module):
    def __init__(
        self, in_channels, out_channels, activation
    ):  # res = 2, ..., resolution_log2
        super().__init__()
        self.conv0 = Conv2dLayer(
            in_channels=in_channels,
            out_channels=in_channels,
            kernel_size=3,
            activation=activation,
        )
        self.conv1 = Conv2dLayer(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=3,
            down=2,
            activation=activation,
        )
        self.skip = Conv2dLayer(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=1,
            down=2,
            bias=False,
        )

    def forward(self, x):
        skip = self.skip(x, gain=np.sqrt(0.5))
        x = self.conv0(x)
        x = self.conv1(x, gain=np.sqrt(0.5))
        out = skip + x

        return out


def nf(stage, channel_base=32768, channel_decay=1.0, channel_max=512):
    NF = {512: 64, 256: 128, 128: 256, 64: 512, 32: 512, 16: 512, 8: 512, 4: 512}
    return NF[2**stage]


class Mlp(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        drop=0.0,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = FullyConnectedLayer(
            in_features=in_features, out_features=hidden_features, activation="lrelu"
        )
        self.fc2 = FullyConnectedLayer(
            in_features=hidden_features, out_features=out_features
        )

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x


def window_partition(x, window_size):
    """
    Args:
        x: (B, H, W, C)
        window_size (int): window size
    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = (
        x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    )
    return windows


def window_reverse(windows, window_size: int, H: int, W: int):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size (int): Window size
        H (int): Height of image
        W (int): Width of image
    Returns:
        x: (B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    # B = windows.shape[0] / (H * W / window_size / window_size)
    x = windows.view(
        B, H // window_size, W // window_size, window_size, window_size, -1
    )
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x


class Conv2dLayerPartial(nn.Module):
    def __init__(
        self,
        in_channels,  # Number of input channels.
        out_channels,  # Number of output channels.
        kernel_size,  # Width and height of the convolution kernel.
        bias=True,  # Apply additive bias before the activation function?
        activation="linear",  # Activation function: 'relu', 'lrelu', etc.
        up=1,  # Integer upsampling factor.
        down=1,  # Integer downsampling factor.
        resample_filter=None,  # Low-pass filter to apply when resampling activations.
        conv_clamp=None,  # Clamp the output to +-X, None = disable clamping.
        trainable=True,  # Update the weights of this layer during training?
    ):
        if resample_filter is None:
            resample_filter = [1, 3, 3, 1]
        super().__init__()
        self.conv = Conv2dLayer(
            in_channels,
            out_channels,
            kernel_size,
            bias,
            activation,
            up,
            down,
            resample_filter,
            conv_clamp,
            trainable,
        )

        self.weight_maskUpdater = torch.ones(1, 1, kernel_size, kernel_size)
        self.slide_winsize = kernel_size**2
        self.stride = down
        self.padding = kernel_size // 2 if kernel_size % 2 == 1 else 0

    def forward(self, x, mask=None):
        if mask is not None:
            with torch.no_grad():
                if self.weight_maskUpdater.type() != x.type():
                    self.weight_maskUpdater = self.weight_maskUpdater.to(x)
                update_mask = F.conv2d(
                    mask,
                    self.weight_maskUpdater,
                    bias=None,
                    stride=self.stride,
                    padding=self.padding,
                )
                mask_ratio = self.slide_winsize / (update_mask + 1e-8)
                update_mask = torch.clamp(update_mask, 0, 1)  # 0 or 1
                mask_ratio = torch.mul(mask_ratio, update_mask)
            x = self.conv(x)
            x = torch.mul(x, mask_ratio)
            return x, update_mask
        else:
            x = self.conv(x)
            return x, None


class WindowAttention(nn.Module):
    r"""Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.
    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
    """

    def __init__(
        self,
        dim,
        window_size,
        num_heads,
        down_ratio=1,
        qkv_bias=True,
        qk_scale=None,
        attn_drop=0.0,
        proj_drop=0.0,
    ):
        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wh, Ww
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        self.q = FullyConnectedLayer(in_features=dim, out_features=dim)
        self.k = FullyConnectedLayer(in_features=dim, out_features=dim)
        self.v = FullyConnectedLayer(in_features=dim, out_features=dim)
        self.proj = FullyConnectedLayer(in_features=dim, out_features=dim)

        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x, mask_windows=None, mask=None):
        """
        Args:
            x: input features with shape of (num_windows*B, N, C)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        B_, N, C = x.shape
        norm_x = F.normalize(x, p=2.0, dim=-1)
        q = (
            self.q(norm_x)
            .reshape(B_, N, self.num_heads, C // self.num_heads)
            .permute(0, 2, 1, 3)
        )
        k = (
            self.k(norm_x)
            .view(B_, -1, self.num_heads, C // self.num_heads)
            .permute(0, 2, 3, 1)
        )
        v = (
            self.v(x)
            .view(B_, -1, self.num_heads, C // self.num_heads)
            .permute(0, 2, 1, 3)
        )

        attn = (q @ k) * self.scale

        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(
                1
            ).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, N, N)

        if mask_windows is not None:
            attn_mask_windows = mask_windows.squeeze(-1).unsqueeze(1).unsqueeze(1)
            attn = attn + attn_mask_windows.masked_fill(
                attn_mask_windows == 0, float(-100.0)
            ).masked_fill(attn_mask_windows == 1, float(0.0))
            with torch.no_grad():
                mask_windows = torch.clamp(
                    torch.sum(mask_windows, dim=1, keepdim=True), 0, 1
                ).repeat(1, N, 1)

        attn = self.softmax(attn)

        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
        x = self.proj(x)
        return x, mask_windows


class SwinTransformerBlock(nn.Module):
    r"""Swin Transformer Block.
    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resulotion.
        num_heads (int): Number of attention heads.
        window_size (int): Window size.
        shift_size (int): Shift size for SW-MSA.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float, optional): Stochastic depth rate. Default: 0.0
        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(
        self,
        dim,
        input_resolution,
        num_heads,
        down_ratio=1,
        window_size=7,
        shift_size=0,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        if min(self.input_resolution) <= self.window_size:
            # if window size is larger than input resolution, we don't partition windows
            self.shift_size = 0
            self.window_size = min(self.input_resolution)
        assert (
            0 <= self.shift_size < self.window_size
        ), "shift_size must in 0-window_size"

        if self.shift_size > 0:
            down_ratio = 1
        self.attn = WindowAttention(
            dim,
            window_size=to_2tuple(self.window_size),
            num_heads=num_heads,
            down_ratio=down_ratio,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            attn_drop=attn_drop,
            proj_drop=drop,
        )

        self.fuse = FullyConnectedLayer(
            in_features=dim * 2, out_features=dim, activation="lrelu"
        )

        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

        if self.shift_size > 0:
            attn_mask = self.calculate_mask(self.input_resolution)
        else:
            attn_mask = None

        self.register_buffer("attn_mask", attn_mask)

    def calculate_mask(self, x_size):
        # calculate attention mask for SW-MSA
        H, W = x_size
        img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
        h_slices = (
            slice(0, -self.window_size),
            slice(-self.window_size, -self.shift_size),
            slice(-self.shift_size, None),
        )
        w_slices = (
            slice(0, -self.window_size),
            slice(-self.window_size, -self.shift_size),
            slice(-self.shift_size, None),
        )
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                img_mask[:, h, w, :] = cnt
                cnt += 1

        mask_windows = window_partition(
            img_mask, self.window_size
        )  # nW, window_size, window_size, 1
        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
            attn_mask == 0, float(0.0)
        )

        return attn_mask

    def forward(self, x, x_size, mask=None):
        # H, W = self.input_resolution
        H, W = x_size
        B, _, C = x.shape
        # assert L == H * W, "input feature has wrong size"

        shortcut = x
        x = x.view(B, H, W, C)
        if mask is not None:
            mask = mask.view(B, H, W, 1)

        # cyclic shift
        if self.shift_size > 0:
            shifted_x = torch.roll(
                x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)
            )
            if mask is not None:
                shifted_mask = torch.roll(
                    mask, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)
                )
        else:
            shifted_x = x
            if mask is not None:
                shifted_mask = mask

        # partition windows
        x_windows = window_partition(
            shifted_x, self.window_size
        )  # nW*B, window_size, window_size, C
        x_windows = x_windows.view(
            -1, self.window_size * self.window_size, C
        )  # nW*B, window_size*window_size, C
        if mask is not None:
            mask_windows = window_partition(shifted_mask, self.window_size)
            mask_windows = mask_windows.view(-1, self.window_size * self.window_size, 1)
        else:
            mask_windows = None

        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
        if self.input_resolution == x_size:
            attn_windows, mask_windows = self.attn(
                x_windows, mask_windows, mask=self.attn_mask
            )  # nW*B, window_size*window_size, C
        else:
            attn_windows, mask_windows = self.attn(
                x_windows, mask_windows, mask=self.calculate_mask(x_size).to(x.device)
            )  # nW*B, window_size*window_size, C

        # merge windows
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
        if mask is not None:
            mask_windows = mask_windows.view(-1, self.window_size, self.window_size, 1)
            shifted_mask = window_reverse(mask_windows, self.window_size, H, W)

        # reverse cyclic shift
        if self.shift_size > 0:
            x = torch.roll(
                shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)
            )
            if mask is not None:
                mask = torch.roll(
                    shifted_mask, shifts=(self.shift_size, self.shift_size), dims=(1, 2)
                )
        else:
            x = shifted_x
            if mask is not None:
                mask = shifted_mask
        x = x.view(B, H * W, C)
        if mask is not None:
            mask = mask.view(B, H * W, 1)

        # FFN
        x = self.fuse(torch.cat([shortcut, x], dim=-1))
        x = self.mlp(x)

        return x, mask


class PatchMerging(nn.Module):
    def __init__(self, in_channels, out_channels, down=2):
        super().__init__()
        self.conv = Conv2dLayerPartial(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=3,
            activation="lrelu",
            down=down,
        )
        self.down = down

    def forward(self, x, x_size, mask=None):
        x = token2feature(x, x_size)
        if mask is not None:
            mask = token2feature(mask, x_size)
        x, mask = self.conv(x, mask)
        if self.down != 1:
            ratio = 1 / self.down
            x_size = (int(x_size[0] * ratio), int(x_size[1] * ratio))
        x = feature2token(x)
        if mask is not None:
            mask = feature2token(mask)
        return x, x_size, mask


class PatchUpsampling(nn.Module):
    def __init__(self, in_channels, out_channels, up=2):
        super().__init__()
        self.conv = Conv2dLayerPartial(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=3,
            activation="lrelu",
            up=up,
        )
        self.up = up

    def forward(self, x, x_size, mask=None):
        x = token2feature(x, x_size)
        if mask is not None:
            mask = token2feature(mask, x_size)
        x, mask = self.conv(x, mask)
        if self.up != 1:
            x_size = (int(x_size[0] * self.up), int(x_size[1] * self.up))
        x = feature2token(x)
        if mask is not None:
            mask = feature2token(mask)
        return x, x_size, mask


class BasicLayer(nn.Module):
    """A basic Swin Transformer layer for one stage.
    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
    """

    def __init__(
        self,
        dim,
        input_resolution,
        depth,
        num_heads,
        window_size,
        down_ratio=1,
        mlp_ratio=2.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        downsample=None,
        use_checkpoint=False,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth
        self.use_checkpoint = use_checkpoint

        # patch merging layer
        if downsample is not None:
            # self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
            self.downsample = downsample
        else:
            self.downsample = None

        # build blocks
        self.blocks = nn.ModuleList(
            [
                SwinTransformerBlock(
                    dim=dim,
                    input_resolution=input_resolution,
                    num_heads=num_heads,
                    down_ratio=down_ratio,
                    window_size=window_size,
                    shift_size=0 if (i % 2 == 0) else window_size // 2,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop,
                    attn_drop=attn_drop,
                    drop_path=drop_path[i]
                    if isinstance(drop_path, list)
                    else drop_path,
                    norm_layer=norm_layer,
                )
                for i in range(depth)
            ]
        )

        self.conv = Conv2dLayerPartial(
            in_channels=dim, out_channels=dim, kernel_size=3, activation="lrelu"
        )

    def forward(self, x, x_size, mask=None):
        if self.downsample is not None:
            x, x_size, mask = self.downsample(x, x_size, mask)
        identity = x
        for blk in self.blocks:
            if self.use_checkpoint:
                x, mask = checkpoint.checkpoint(blk, x, x_size, mask)
            else:
                x, mask = blk(x, x_size, mask)
        if mask is not None:
            mask = token2feature(mask, x_size)
        x, mask = self.conv(token2feature(x, x_size), mask)
        x = feature2token(x) + identity
        if mask is not None:
            mask = feature2token(mask)
        return x, x_size, mask


class ToToken(nn.Module):
    def __init__(self, in_channels=3, dim=128, kernel_size=5, stride=1):
        super().__init__()

        self.proj = Conv2dLayerPartial(
            in_channels=in_channels,
            out_channels=dim,
            kernel_size=kernel_size,
            activation="lrelu",
        )

    def forward(self, x, mask):
        x, mask = self.proj(x, mask)

        return x, mask


class EncFromRGB(nn.Module):
    def __init__(
        self, in_channels, out_channels, activation
    ):  # res = 2, ..., resolution_log2
        super().__init__()
        self.conv0 = Conv2dLayer(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=1,
            activation=activation,
        )
        self.conv1 = Conv2dLayer(
            in_channels=out_channels,
            out_channels=out_channels,
            kernel_size=3,
            activation=activation,
        )

    def forward(self, x):
        x = self.conv0(x)
        x = self.conv1(x)

        return x


class ConvBlockDown(nn.Module):
    def __init__(
        self, in_channels, out_channels, activation
    ):  # res = 2, ..., resolution_log
        super().__init__()

        self.conv0 = Conv2dLayer(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=3,
            activation=activation,
            down=2,
        )
        self.conv1 = Conv2dLayer(
            in_channels=out_channels,
            out_channels=out_channels,
            kernel_size=3,
            activation=activation,
        )

    def forward(self, x):
        x = self.conv0(x)
        x = self.conv1(x)

        return x


def token2feature(x, x_size):
    B, _, C = x.shape
    h, w = x_size
    x = x.permute(0, 2, 1).reshape(B, C, h, w)
    return x


def feature2token(x):
    B, C, _, _ = x.shape
    x = x.view(B, C, -1).transpose(1, 2)
    return x


class Encoder(nn.Module):
    def __init__(
        self,
        res_log2,
        img_channels,
        activation,
        patch_size=5,
        channels=16,
        drop_path_rate=0.1,
    ):
        super().__init__()

        self.resolution = []

        for i in range(res_log2, 3, -1):  # from input size to 16x16
            res = 2**i
            self.resolution.append(res)
            if i == res_log2:
                block = EncFromRGB(img_channels * 2 + 1, nf(i), activation)
            else:
                block = ConvBlockDown(nf(i + 1), nf(i), activation)
            setattr(self, "EncConv_Block_%dx%d" % (res, res), block)

    def forward(self, x):
        out = {}
        for res in self.resolution:
            res_log2 = int(np.log2(res))
            x = getattr(self, "EncConv_Block_%dx%d" % (res, res))(x)
            out[res_log2] = x

        return out


class ToStyle(nn.Module):
    def __init__(self, in_channels, out_channels, activation, drop_rate):
        super().__init__()
        self.conv = nn.Sequential(
            Conv2dLayer(
                in_channels=in_channels,
                out_channels=in_channels,
                kernel_size=3,
                activation=activation,
                down=2,
            ),
            Conv2dLayer(
                in_channels=in_channels,
                out_channels=in_channels,
                kernel_size=3,
                activation=activation,
                down=2,
            ),
            Conv2dLayer(
                in_channels=in_channels,
                out_channels=in_channels,
                kernel_size=3,
                activation=activation,
                down=2,
            ),
        )

        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = FullyConnectedLayer(
            in_features=in_channels, out_features=out_channels, activation=activation
        )
        # self.dropout = nn.Dropout(drop_rate)

    def forward(self, x):
        x = self.conv(x)
        x = self.pool(x)
        x = self.fc(x.flatten(start_dim=1))
        # x = self.dropout(x)

        return x


class DecBlockFirstV2(nn.Module):
    def __init__(
        self,
        res,
        in_channels,
        out_channels,
        activation,
        style_dim,
        use_noise,
        demodulate,
        img_channels,
    ):
        super().__init__()
        self.res = res

        self.conv0 = Conv2dLayer(
            in_channels=in_channels,
            out_channels=in_channels,
            kernel_size=3,
            activation=activation,
        )
        self.conv1 = StyleConv(
            in_channels=in_channels,
            out_channels=out_channels,
            style_dim=style_dim,
            resolution=2**res,
            kernel_size=3,
            use_noise=use_noise,
            activation=activation,
            demodulate=demodulate,
        )
        self.toRGB = ToRGB(
            in_channels=out_channels,
            out_channels=img_channels,
            style_dim=style_dim,
            kernel_size=1,
            demodulate=False,
        )

    def forward(self, x, ws, gs, E_features, noise_mode="random"):
        # x = self.fc(x).view(x.shape[0], -1, 4, 4)
        x = self.conv0(x)
        x = x + E_features[self.res]
        style = get_style_code(ws[:, 0], gs)
        x = self.conv1(x, style, noise_mode=noise_mode)
        style = get_style_code(ws[:, 1], gs)
        img = self.toRGB(x, style, skip=None)

        return x, img


class DecBlock(nn.Module):
    def __init__(
        self,
        res,
        in_channels,
        out_channels,
        activation,
        style_dim,
        use_noise,
        demodulate,
        img_channels,
    ):  # res = 4, ..., resolution_log2
        super().__init__()
        self.res = res

        self.conv0 = StyleConv(
            in_channels=in_channels,
            out_channels=out_channels,
            style_dim=style_dim,
            resolution=2**res,
            kernel_size=3,
            up=2,
            use_noise=use_noise,
            activation=activation,
            demodulate=demodulate,
        )
        self.conv1 = StyleConv(
            in_channels=out_channels,
            out_channels=out_channels,
            style_dim=style_dim,
            resolution=2**res,
            kernel_size=3,
            use_noise=use_noise,
            activation=activation,
            demodulate=demodulate,
        )
        self.toRGB = ToRGB(
            in_channels=out_channels,
            out_channels=img_channels,
            style_dim=style_dim,
            kernel_size=1,
            demodulate=False,
        )

    def forward(self, x, img, ws, gs, E_features, noise_mode="random"):
        style = get_style_code(ws[:, self.res * 2 - 9], gs)
        x = self.conv0(x, style, noise_mode=noise_mode)
        x = x + E_features[self.res]
        style = get_style_code(ws[:, self.res * 2 - 8], gs)
        x = self.conv1(x, style, noise_mode=noise_mode)
        style = get_style_code(ws[:, self.res * 2 - 7], gs)
        img = self.toRGB(x, style, skip=img)

        return x, img


class Decoder(nn.Module):
    def __init__(
        self, res_log2, activation, style_dim, use_noise, demodulate, img_channels
    ):
        super().__init__()
        self.Dec_16x16 = DecBlockFirstV2(
            4, nf(4), nf(4), activation, style_dim, use_noise, demodulate, img_channels
        )
        for res in range(5, res_log2 + 1):
            setattr(
                self,
                "Dec_%dx%d" % (2**res, 2**res),
                DecBlock(
                    res,
                    nf(res - 1),
                    nf(res),
                    activation,
                    style_dim,
                    use_noise,
                    demodulate,
                    img_channels,
                ),
            )
        self.res_log2 = res_log2

    def forward(self, x, ws, gs, E_features, noise_mode="random"):
        x, img = self.Dec_16x16(x, ws, gs, E_features, noise_mode=noise_mode)
        for res in range(5, self.res_log2 + 1):
            block = getattr(self, "Dec_%dx%d" % (2**res, 2**res))
            x, img = block(x, img, ws, gs, E_features, noise_mode=noise_mode)

        return img


class DecStyleBlock(nn.Module):
    def __init__(
        self,
        res,
        in_channels,
        out_channels,
        activation,
        style_dim,
        use_noise,
        demodulate,
        img_channels,
    ):
        super().__init__()
        self.res = res

        self.conv0 = StyleConv(
            in_channels=in_channels,
            out_channels=out_channels,
            style_dim=style_dim,
            resolution=2**res,
            kernel_size=3,
            up=2,
            use_noise=use_noise,
            activation=activation,
            demodulate=demodulate,
        )
        self.conv1 = StyleConv(
            in_channels=out_channels,
            out_channels=out_channels,
            style_dim=style_dim,
            resolution=2**res,
            kernel_size=3,
            use_noise=use_noise,
            activation=activation,
            demodulate=demodulate,
        )
        self.toRGB = ToRGB(
            in_channels=out_channels,
            out_channels=img_channels,
            style_dim=style_dim,
            kernel_size=1,
            demodulate=False,
        )

    def forward(self, x, img, style, skip, noise_mode="random"):
        x = self.conv0(x, style, noise_mode=noise_mode)
        x = x + skip
        x = self.conv1(x, style, noise_mode=noise_mode)
        img = self.toRGB(x, style, skip=img)

        return x, img


class FirstStage(nn.Module):
    def __init__(
        self,
        img_channels,
        img_resolution=256,
        dim=180,
        w_dim=512,
        use_noise=False,
        demodulate=True,
        activation="lrelu",
    ):
        super().__init__()
        res = 64

        self.conv_first = Conv2dLayerPartial(
            in_channels=img_channels + 1,
            out_channels=dim,
            kernel_size=3,
            activation=activation,
        )
        self.enc_conv = nn.ModuleList()
        down_time = int(np.log2(img_resolution // res))
        # 根据图片尺寸构建 swim transformer 的层数
        for i in range(down_time):  # from input size to 64
            self.enc_conv.append(
                Conv2dLayerPartial(
                    in_channels=dim,
                    out_channels=dim,
                    kernel_size=3,
                    down=2,
                    activation=activation,
                )
            )

        # from 64 -> 16 -> 64
        depths = [2, 3, 4, 3, 2]
        ratios = [1, 1 / 2, 1 / 2, 2, 2]
        num_heads = 6
        window_sizes = [8, 16, 16, 16, 8]
        drop_path_rate = 0.1
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]

        self.tran = nn.ModuleList()
        for i, depth in enumerate(depths):
            res = int(res * ratios[i])
            if ratios[i] < 1:
                merge = PatchMerging(dim, dim, down=int(1 / ratios[i]))
            elif ratios[i] > 1:
                merge = PatchUpsampling(dim, dim, up=ratios[i])
            else:
                merge = None
            self.tran.append(
                BasicLayer(
                    dim=dim,
                    input_resolution=[res, res],
                    depth=depth,
                    num_heads=num_heads,
                    window_size=window_sizes[i],
                    drop_path=dpr[sum(depths[:i]) : sum(depths[: i + 1])],
                    downsample=merge,
                )
            )

        # global style
        down_conv = []
        for i in range(int(np.log2(16))):
            down_conv.append(
                Conv2dLayer(
                    in_channels=dim,
                    out_channels=dim,
                    kernel_size=3,
                    down=2,
                    activation=activation,
                )
            )
        down_conv.append(nn.AdaptiveAvgPool2d((1, 1)))
        self.down_conv = nn.Sequential(*down_conv)
        self.to_style = FullyConnectedLayer(
            in_features=dim, out_features=dim * 2, activation=activation
        )
        self.ws_style = FullyConnectedLayer(
            in_features=w_dim, out_features=dim, activation=activation
        )
        self.to_square = FullyConnectedLayer(
            in_features=dim, out_features=16 * 16, activation=activation
        )

        style_dim = dim * 3
        self.dec_conv = nn.ModuleList()
        for i in range(down_time):  # from 64 to input size
            res = res * 2
            self.dec_conv.append(
                DecStyleBlock(
                    res,
                    dim,
                    dim,
                    activation,
                    style_dim,
                    use_noise,
                    demodulate,
                    img_channels,
                )
            )

    def forward(self, images_in, masks_in, ws, noise_mode="random"):
        x = torch.cat([masks_in - 0.5, images_in * masks_in], dim=1)

        skips = []
        x, mask = self.conv_first(x, masks_in)  # input size
        skips.append(x)
        for i, block in enumerate(self.enc_conv):  # input size to 64
            x, mask = block(x, mask)
            if i != len(self.enc_conv) - 1:
                skips.append(x)

        x_size = x.size()[-2:]
        x = feature2token(x)
        mask = feature2token(mask)
        mid = len(self.tran) // 2
        for i, block in enumerate(self.tran):  # 64 to 16
            if i < mid:
                x, x_size, mask = block(x, x_size, mask)
                skips.append(x)
            elif i > mid:
                x, x_size, mask = block(x, x_size, None)
                x = x + skips[mid - i]
            else:
                x, x_size, mask = block(x, x_size, None)

                mul_map = torch.ones_like(x) * 0.5
                mul_map = F.dropout(mul_map, training=True).to(x.device)
                ws = self.ws_style(ws[:, -1]).to(x.device)
                add_n = self.to_square(ws).unsqueeze(1).to(x.device)
                add_n = (
                    F.interpolate(
                        add_n, size=x.size(1), mode="linear", align_corners=False
                    )
                    .squeeze(1)
                    .unsqueeze(-1)
                ).to(x.device)
                x = x * mul_map + add_n * (1 - mul_map)
                gs = self.to_style(
                    self.down_conv(token2feature(x, x_size)).flatten(start_dim=1)
                ).to(x.device)
                style = torch.cat([gs, ws], dim=1)

        x = token2feature(x, x_size).contiguous()
        img = None
        for i, block in enumerate(self.dec_conv):
            x, img = block(
                x, img, style, skips[len(self.dec_conv) - i - 1], noise_mode=noise_mode
            )

        # ensemble
        img = img * (1 - masks_in) + images_in * masks_in

        return img


class SynthesisNet(nn.Module):
    def __init__(
        self,
        w_dim,  # Intermediate latent (W) dimensionality.
        img_resolution,  # Output image resolution.
        img_channels=3,  # Number of color channels.
        channel_base=32768,  # Overall multiplier for the number of channels.
        channel_decay=1.0,
        channel_max=512,  # Maximum number of channels in any layer.
        activation="lrelu",  # Activation function: 'relu', 'lrelu', etc.
        drop_rate=0.5,
        use_noise=False,
        demodulate=True,
    ):
        super().__init__()
        resolution_log2 = int(np.log2(img_resolution))
        assert img_resolution == 2**resolution_log2 and img_resolution >= 4

        self.num_layers = resolution_log2 * 2 - 3 * 2
        self.img_resolution = img_resolution
        self.resolution_log2 = resolution_log2

        # first stage
        self.first_stage = FirstStage(
            img_channels,
            img_resolution=img_resolution,
            w_dim=w_dim,
            use_noise=False,
            demodulate=demodulate,
        )

        # second stage
        self.enc = Encoder(
            resolution_log2, img_channels, activation, patch_size=5, channels=16
        )
        self.to_square = FullyConnectedLayer(
            in_features=w_dim, out_features=16 * 16, activation=activation
        )
        self.to_style = ToStyle(
            in_channels=nf(4),
            out_channels=nf(2) * 2,
            activation=activation,
            drop_rate=drop_rate,
        )
        style_dim = w_dim + nf(2) * 2
        self.dec = Decoder(
            resolution_log2, activation, style_dim, use_noise, demodulate, img_channels
        )

    def forward(self, images_in, masks_in, ws, noise_mode="random", return_stg1=False):
        out_stg1 = self.first_stage(images_in, masks_in, ws, noise_mode=noise_mode)

        # encoder
        x = images_in * masks_in + out_stg1 * (1 - masks_in)
        x = torch.cat([masks_in - 0.5, x, images_in * masks_in], dim=1)
        E_features = self.enc(x)

        fea_16 = E_features[4].to(x.device)
        mul_map = torch.ones_like(fea_16) * 0.5
        mul_map = F.dropout(mul_map, training=True).to(x.device)
        add_n = self.to_square(ws[:, 0]).view(-1, 16, 16).unsqueeze(1)
        add_n = F.interpolate(
            add_n, size=fea_16.size()[-2:], mode="bilinear", align_corners=False
        ).to(x.device)
        fea_16 = fea_16 * mul_map + add_n * (1 - mul_map)
        E_features[4] = fea_16

        # style
        gs = self.to_style(fea_16).to(x.device)

        # decoder
        img = self.dec(fea_16, ws, gs, E_features, noise_mode=noise_mode).to(x.device)

        # ensemble
        img = img * (1 - masks_in) + images_in * masks_in

        if not return_stg1:
            return img
        else:
            return img, out_stg1


class Generator(nn.Module):
    def __init__(
        self,
        z_dim,  # Input latent (Z) dimensionality, 0 = no latent.
        c_dim,  # Conditioning label (C) dimensionality, 0 = no label.
        w_dim,  # Intermediate latent (W) dimensionality.
        img_resolution,  # resolution of generated image
        img_channels,  # Number of input color channels.
        synthesis_kwargs=None,  # Arguments for SynthesisNetwork.
        mapping_kwargs=None,  # Arguments for MappingNetwork.
    ):
        if mapping_kwargs is None:
            mapping_kwargs = {}
        if synthesis_kwargs is None:
            synthesis_kwargs = {}
        super().__init__()
        self.z_dim = z_dim
        self.c_dim = c_dim
        self.w_dim = w_dim
        self.img_resolution = img_resolution
        self.img_channels = img_channels

        self.synthesis = SynthesisNet(
            w_dim=w_dim,
            img_resolution=img_resolution,
            img_channels=img_channels,
            **synthesis_kwargs,
        )
        self.mapping = MappingNet(
            z_dim=z_dim,
            c_dim=c_dim,
            w_dim=w_dim,
            num_ws=self.synthesis.num_layers,
            **mapping_kwargs,
        )

    def forward(
        self,
        images_in,
        masks_in,
        z,
        c,
        truncation_psi=1,
        truncation_cutoff=None,
        skip_w_avg_update=False,
        noise_mode="none",
        return_stg1=False,
    ):
        ws = self.mapping(
            z,
            c,
            truncation_psi=truncation_psi,
            truncation_cutoff=truncation_cutoff,
            skip_w_avg_update=skip_w_avg_update,
        )
        img = self.synthesis(images_in, masks_in, ws, noise_mode=noise_mode)
        return img


class MAT(nn.Module):
    def __init__(self, state_dict):
        super(MAT, self).__init__()
        self.model_arch = "MAT"
        self.sub_type = "Inpaint"
        self.in_nc = 3
        self.out_nc = 3
        self.scale = 1

        self.supports_fp16 = False
        self.supports_bf16 = True

        self.min_size = 512
        self.pad_mod = 512
        self.pad_to_square = True

        seed = 240  # pick up a random number
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)

        self.model = Generator(
            z_dim=512, c_dim=0, w_dim=512, img_resolution=512, img_channels=3
        )
        self.z = torch.from_numpy(np.random.randn(1, self.model.z_dim))  # [1., 512]
        self.label = torch.zeros([1, self.model.c_dim])
        self.state = {
            k.replace("synthesis", "model.synthesis").replace(
                "mapping", "model.mapping"
            ): v
            for k, v in state_dict.items()
        }
        self.load_state_dict(self.state, strict=False)

    def forward(self, image, mask):
        """Input images and output images have same size
        images: [H, W, C] RGB
        masks: [H, W] mask area == 255
        return: BGR IMAGE
        """

        image = image * 2 - 1  # [0, 1] -> [-1, 1]
        mask = 1 - mask

        output = self.model(
            image, mask, self.z, self.label, truncation_psi=1, noise_mode="none"
        )

        return output * 0.5 + 0.5

nodes/impl/pytorch/architecture/MAT.py
vladmandic/sd-extension-chainner
import math

import torch.nn as nn


class CA_layer(nn.Module):
    def __init__(self, channel, reduction=16):
        super(CA_layer, self).__init__()
        # global average pooling
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Conv2d(channel, channel // reduction, kernel_size=(1, 1), bias=False),
            nn.GELU(),
            nn.Conv2d(channel // reduction, channel, kernel_size=(1, 1), bias=False),
            # nn.Sigmoid()
        )

    def forward(self, x):
        y = self.fc(self.gap(x))
        return x * y.expand_as(x)


class Simple_CA_layer(nn.Module):
    def __init__(self, channel):
        super(Simple_CA_layer, self).__init__()
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Conv2d(
            in_channels=channel,
            out_channels=channel,
            kernel_size=1,
            padding=0,
            stride=1,
            groups=1,
            bias=True,
        )

    def forward(self, x):
        return x * self.fc(self.gap(x))


class ECA_layer(nn.Module):
    """Constructs a ECA module.
    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """

    def __init__(self, channel):
        super(ECA_layer, self).__init__()

        b = 1
        gamma = 2
        k_size = int(abs(math.log(channel, 2) + b) / gamma)
        k_size = k_size if k_size % 2 else k_size + 1
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(
            1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False
        )
        # self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # x: input features with shape [b, c, h, w]
        # b, c, h, w = x.size()

        # feature descriptor on the global spatial information
        y = self.avg_pool(x)

        # Two different branches of ECA module
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Multi-scale information fusion
        # y = self.sigmoid(y)

        return x * y.expand_as(x)


class ECA_MaxPool_layer(nn.Module):
    """Constructs a ECA module.
    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """

    def __init__(self, channel):
        super(ECA_MaxPool_layer, self).__init__()

        b = 1
        gamma = 2
        k_size = int(abs(math.log(channel, 2) + b) / gamma)
        k_size = k_size if k_size % 2 else k_size + 1
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.conv = nn.Conv1d(
            1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False
        )
        # self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # x: input features with shape [b, c, h, w]
        # b, c, h, w = x.size()

        # feature descriptor on the global spatial information
        y = self.max_pool(x)

        # Two different branches of ECA module
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Multi-scale information fusion
        # y = self.sigmoid(y)

        return x * y.expand_as(x)

nodes/impl/pytorch/architecture/OmniSR/ChannelAttention.py
vladmandic/sd-extension-chainner
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: OSA.py
# Created Date: Tuesday April 28th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified:  Sunday, 23rd April 2023 3:07:42 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2020 Shanghai Jiao Tong University
#############################################################

import torch
import torch.nn.functional as F
from einops import rearrange, repeat
from einops.layers.torch import Rearrange, Reduce
from torch import einsum, nn

from .layernorm import LayerNorm2d

# helpers


def exists(val):
    return val is not None


def default(val, d):
    return val if exists(val) else d


def cast_tuple(val, length=1):
    return val if isinstance(val, tuple) else ((val,) * length)


# helper classes


class PreNormResidual(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x):
        return self.fn(self.norm(x)) + x


class Conv_PreNormResidual(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = LayerNorm2d(dim)
        self.fn = fn

    def forward(self, x):
        return self.fn(self.norm(x)) + x


class FeedForward(nn.Module):
    def __init__(self, dim, mult=2, dropout=0.0):
        super().__init__()
        inner_dim = int(dim * mult)
        self.net = nn.Sequential(
            nn.Linear(dim, inner_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


class Conv_FeedForward(nn.Module):
    def __init__(self, dim, mult=2, dropout=0.0):
        super().__init__()
        inner_dim = int(dim * mult)
        self.net = nn.Sequential(
            nn.Conv2d(dim, inner_dim, 1, 1, 0),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Conv2d(inner_dim, dim, 1, 1, 0),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


class Gated_Conv_FeedForward(nn.Module):
    def __init__(self, dim, mult=1, bias=False, dropout=0.0):
        super().__init__()

        hidden_features = int(dim * mult)

        self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)

        self.dwconv = nn.Conv2d(
            hidden_features * 2,
            hidden_features * 2,
            kernel_size=3,
            stride=1,
            padding=1,
            groups=hidden_features * 2,
            bias=bias,
        )

        self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)

    def forward(self, x):
        x = self.project_in(x)
        x1, x2 = self.dwconv(x).chunk(2, dim=1)
        x = F.gelu(x1) * x2
        x = self.project_out(x)
        return x


# MBConv


class SqueezeExcitation(nn.Module):
    def __init__(self, dim, shrinkage_rate=0.25):
        super().__init__()
        hidden_dim = int(dim * shrinkage_rate)

        self.gate = nn.Sequential(
            Reduce("b c h w -> b c", "mean"),
            nn.Linear(dim, hidden_dim, bias=False),
            nn.SiLU(),
            nn.Linear(hidden_dim, dim, bias=False),
            nn.Sigmoid(),
            Rearrange("b c -> b c 1 1"),
        )

    def forward(self, x):
        return x * self.gate(x)


class MBConvResidual(nn.Module):
    def __init__(self, fn, dropout=0.0):
        super().__init__()
        self.fn = fn
        self.dropsample = Dropsample(dropout)

    def forward(self, x):
        out = self.fn(x)
        out = self.dropsample(out)
        return out + x


class Dropsample(nn.Module):
    def __init__(self, prob=0):
        super().__init__()
        self.prob = prob

    def forward(self, x):
        device = x.device

        if self.prob == 0.0 or (not self.training):
            return x

        keep_mask = (
            torch.FloatTensor((x.shape[0], 1, 1, 1), device=device).uniform_()
            > self.prob
        )
        return x * keep_mask / (1 - self.prob)


def MBConv(
    dim_in, dim_out, *, downsample, expansion_rate=4, shrinkage_rate=0.25, dropout=0.0
):
    hidden_dim = int(expansion_rate * dim_out)
    stride = 2 if downsample else 1

    net = nn.Sequential(
        nn.Conv2d(dim_in, hidden_dim, 1),
        # nn.BatchNorm2d(hidden_dim),
        nn.GELU(),
        nn.Conv2d(
            hidden_dim, hidden_dim, 3, stride=stride, padding=1, groups=hidden_dim
        ),
        # nn.BatchNorm2d(hidden_dim),
        nn.GELU(),
        SqueezeExcitation(hidden_dim, shrinkage_rate=shrinkage_rate),
        nn.Conv2d(hidden_dim, dim_out, 1),
        # nn.BatchNorm2d(dim_out)
    )

    if dim_in == dim_out and not downsample:
        net = MBConvResidual(net, dropout=dropout)

    return net


# attention related classes
class Attention(nn.Module):
    def __init__(
        self,
        dim,
        dim_head=32,
        dropout=0.0,
        window_size=7,
        with_pe=True,
    ):
        super().__init__()
        assert (
            dim % dim_head
        ) == 0, "dimension should be divisible by dimension per head"

        self.heads = dim // dim_head
        self.scale = dim_head**-0.5
        self.with_pe = with_pe

        self.to_qkv = nn.Linear(dim, dim * 3, bias=False)

        self.attend = nn.Sequential(nn.Softmax(dim=-1), nn.Dropout(dropout))

        self.to_out = nn.Sequential(
            nn.Linear(dim, dim, bias=False), nn.Dropout(dropout)
        )

        # relative positional bias
        if self.with_pe:
            self.rel_pos_bias = nn.Embedding((2 * window_size - 1) ** 2, self.heads)

            pos = torch.arange(window_size)
            grid = torch.stack(torch.meshgrid(pos, pos))
            grid = rearrange(grid, "c i j -> (i j) c")
            rel_pos = rearrange(grid, "i ... -> i 1 ...") - rearrange(
                grid, "j ... -> 1 j ..."
            )
            rel_pos += window_size - 1
            rel_pos_indices = (rel_pos * torch.tensor([2 * window_size - 1, 1])).sum(
                dim=-1
            )

            self.register_buffer("rel_pos_indices", rel_pos_indices, persistent=False)

    def forward(self, x):
        batch, height, width, window_height, window_width, _, device, h = (
            *x.shape,
            x.device,
            self.heads,
        )

        # flatten

        x = rearrange(x, "b x y w1 w2 d -> (b x y) (w1 w2) d")

        # project for queries, keys, values

        q, k, v = self.to_qkv(x).chunk(3, dim=-1)

        # split heads

        q, k, v = (rearrange(t, "b n (h d ) -> b h n d", h=h) for t in (q, k, v))

        # scale

        q = q * self.scale

        # sim

        sim = einsum("b h i d, b h j d -> b h i j", q, k)

        # add positional bias
        if self.with_pe:
            bias = self.rel_pos_bias(self.rel_pos_indices)
            sim = sim + rearrange(bias, "i j h -> h i j")

        # attention

        attn = self.attend(sim)

        # aggregate

        out = einsum("b h i j, b h j d -> b h i d", attn, v)

        # merge heads

        out = rearrange(
            out, "b h (w1 w2) d -> b w1 w2 (h d)", w1=window_height, w2=window_width
        )

        # combine heads out

        out = self.to_out(out)
        return rearrange(out, "(b x y) ... -> b x y ...", x=height, y=width)


class Block_Attention(nn.Module):
    def __init__(
        self,
        dim,
        dim_head=32,
        bias=False,
        dropout=0.0,
        window_size=7,
        with_pe=True,
    ):
        super().__init__()
        assert (
            dim % dim_head
        ) == 0, "dimension should be divisible by dimension per head"

        self.heads = dim // dim_head
        self.ps = window_size
        self.scale = dim_head**-0.5
        self.with_pe = with_pe

        self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
        self.qkv_dwconv = nn.Conv2d(
            dim * 3,
            dim * 3,
            kernel_size=3,
            stride=1,
            padding=1,
            groups=dim * 3,
            bias=bias,
        )

        self.attend = nn.Sequential(nn.Softmax(dim=-1), nn.Dropout(dropout))

        self.to_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)

    def forward(self, x):
        # project for queries, keys, values
        b, c, h, w = x.shape

        qkv = self.qkv_dwconv(self.qkv(x))
        q, k, v = qkv.chunk(3, dim=1)

        # split heads

        q, k, v = (rearrange(
                t,
                "b (h d) (x w1) (y w2) -> (b x y) h (w1 w2) d",
                h=self.heads,
                w1=self.ps,
                w2=self.ps,
            ) for t in (q, k, v))

        # scale

        q = q * self.scale

        # sim

        sim = einsum("b h i d, b h j d -> b h i j", q, k)

        # attention
        attn = self.attend(sim)

        # aggregate

        out = einsum("b h i j, b h j d -> b h i d", attn, v)

        # merge heads
        out = rearrange(
            out,
            "(b x y) head (w1 w2) d -> b (head d) (x w1) (y w2)",
            x=h // self.ps,
            y=w // self.ps,
            head=self.heads,
            w1=self.ps,
            w2=self.ps,
        )

        out = self.to_out(out)
        return out


class Channel_Attention(nn.Module):
    def __init__(self, dim, heads, bias=False, dropout=0.0, window_size=7):
        super(Channel_Attention, self).__init__()
        self.heads = heads

        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))

        self.ps = window_size

        self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
        self.qkv_dwconv = nn.Conv2d(
            dim * 3,
            dim * 3,
            kernel_size=3,
            stride=1,
            padding=1,
            groups=dim * 3,
            bias=bias,
        )
        self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)

    def forward(self, x):
        b, c, h, w = x.shape

        qkv = self.qkv_dwconv(self.qkv(x))
        qkv = qkv.chunk(3, dim=1)

        q, k, v = (rearrange(
                t,
                "b (head d) (h ph) (w pw) -> b (h w) head d (ph pw)",
                ph=self.ps,
                pw=self.ps,
                head=self.heads,
            ) for t in qkv)

        q = F.normalize(q, dim=-1)
        k = F.normalize(k, dim=-1)

        attn = (q @ k.transpose(-2, -1)) * self.temperature
        attn = attn.softmax(dim=-1)
        out = attn @ v

        out = rearrange(
            out,
            "b (h w) head d (ph pw) -> b (head d) (h ph) (w pw)",
            h=h // self.ps,
            w=w // self.ps,
            ph=self.ps,
            pw=self.ps,
            head=self.heads,
        )

        out = self.project_out(out)

        return out


class Channel_Attention_grid(nn.Module):
    def __init__(self, dim, heads, bias=False, dropout=0.0, window_size=7):
        super(Channel_Attention_grid, self).__init__()
        self.heads = heads

        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))

        self.ps = window_size

        self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
        self.qkv_dwconv = nn.Conv2d(
            dim * 3,
            dim * 3,
            kernel_size=3,
            stride=1,
            padding=1,
            groups=dim * 3,
            bias=bias,
        )
        self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)

    def forward(self, x):
        b, c, h, w = x.shape

        qkv = self.qkv_dwconv(self.qkv(x))
        qkv = qkv.chunk(3, dim=1)

        q, k, v = (rearrange(
                t,
                "b (head d) (h ph) (w pw) -> b (ph pw) head d (h w)",
                ph=self.ps,
                pw=self.ps,
                head=self.heads,
            ) for t in qkv)

        q = F.normalize(q, dim=-1)
        k = F.normalize(k, dim=-1)

        attn = (q @ k.transpose(-2, -1)) * self.temperature
        attn = attn.softmax(dim=-1)
        out = attn @ v

        out = rearrange(
            out,
            "b (ph pw) head d (h w) -> b (head d) (h ph) (w pw)",
            h=h // self.ps,
            w=w // self.ps,
            ph=self.ps,
            pw=self.ps,
            head=self.heads,
        )

        out = self.project_out(out)

        return out


class OSA_Block(nn.Module):
    def __init__(
        self,
        channel_num=64,
        bias=True,
        ffn_bias=True,
        window_size=8,
        with_pe=False,
        dropout=0.0,
    ):
        super(OSA_Block, self).__init__()

        w = window_size

        self.layer = nn.Sequential(
            MBConv(
                channel_num,
                channel_num,
                downsample=False,
                expansion_rate=1,
                shrinkage_rate=0.25,
            ),
            Rearrange(
                "b d (x w1) (y w2) -> b x y w1 w2 d", w1=w, w2=w
            ),  # block-like attention
            PreNormResidual(
                channel_num,
                Attention(
                    dim=channel_num,
                    dim_head=channel_num // 4,
                    dropout=dropout,
                    window_size=window_size,
                    with_pe=with_pe,
                ),
            ),
            Rearrange("b x y w1 w2 d -> b d (x w1) (y w2)"),
            Conv_PreNormResidual(
                channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout)
            ),
            # channel-like attention
            Conv_PreNormResidual(
                channel_num,
                Channel_Attention(
                    dim=channel_num, heads=4, dropout=dropout, window_size=window_size
                ),
            ),
            Conv_PreNormResidual(
                channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout)
            ),
            Rearrange(
                "b d (w1 x) (w2 y) -> b x y w1 w2 d", w1=w, w2=w
            ),  # grid-like attention
            PreNormResidual(
                channel_num,
                Attention(
                    dim=channel_num,
                    dim_head=channel_num // 4,
                    dropout=dropout,
                    window_size=window_size,
                    with_pe=with_pe,
                ),
            ),
            Rearrange("b x y w1 w2 d -> b d (w1 x) (w2 y)"),
            Conv_PreNormResidual(
                channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout)
            ),
            # channel-like attention
            Conv_PreNormResidual(
                channel_num,
                Channel_Attention_grid(
                    dim=channel_num, heads=4, dropout=dropout, window_size=window_size
                ),
            ),
            Conv_PreNormResidual(
                channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout)
            ),
        )

    def forward(self, x):
        out = self.layer(x)
        return out

nodes/impl/pytorch/architecture/OmniSR/OSA.py
vladmandic/sd-extension-chainner
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: OSAG.py
# Created Date: Tuesday April 28th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified:  Sunday, 23rd April 2023 3:08:49 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2020 Shanghai Jiao Tong University
#############################################################


import torch.nn as nn

from .esa import ESA
from .OSA import OSA_Block


class OSAG(nn.Module):
    def __init__(
        self,
        channel_num=64,
        bias=True,
        block_num=4,
        ffn_bias=False,
        window_size=0,
        pe=False,
    ):
        super(OSAG, self).__init__()

        # print("window_size: %d" % (window_size))
        # print("with_pe", pe)
        # print("ffn_bias: %d" % (ffn_bias))

        # block_script_name = kwargs.get("block_script_name", "OSA")
        # block_class_name = kwargs.get("block_class_name", "OSA_Block")

        # script_name = "." + block_script_name
        # package = __import__(script_name, fromlist=True)
        block_class = OSA_Block  # getattr(package, block_class_name)
        group_list = []
        for _ in range(block_num):
            temp_res = block_class(
                channel_num,
                bias,
                ffn_bias=ffn_bias,
                window_size=window_size,
                with_pe=pe,
            )
            group_list.append(temp_res)
        group_list.append(nn.Conv2d(channel_num, channel_num, 1, 1, 0, bias=bias))
        self.residual_layer = nn.Sequential(*group_list)
        esa_channel = max(channel_num // 4, 16)
        self.esa = ESA(esa_channel, channel_num)

    def forward(self, x):
        out = self.residual_layer(x)
        out = out + x
        return self.esa(out)

nodes/impl/pytorch/architecture/OmniSR/OSAG.py
vladmandic/sd-extension-chainner
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: OmniSR.py
# Created Date: Tuesday April 28th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified:  Sunday, 23rd April 2023 3:06:36 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2020 Shanghai Jiao Tong University
#############################################################

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from .OSAG import OSAG
from .pixelshuffle import pixelshuffle_block


class OmniSR(nn.Module):
    def __init__(
        self,
        state_dict,
        **kwargs,
    ):
        super(OmniSR, self).__init__()
        self.state = state_dict

        bias = True  # Fine to assume this for now
        block_num = 1  # Fine to assume this for now
        ffn_bias = True
        pe = True

        num_feat = state_dict["input.weight"].shape[0] or 64
        num_in_ch = state_dict["input.weight"].shape[1] or 3
        num_out_ch = num_in_ch  # we can just assume this for now. pixelshuffle smh

        pixelshuffle_shape = state_dict["up.0.weight"].shape[0]
        up_scale = math.sqrt(pixelshuffle_shape / num_out_ch)
        if up_scale - int(up_scale) > 0:
            print(
                "out_nc is probably different than in_nc, scale calculation might be wrong"
            )
        up_scale = int(up_scale)
        res_num = 0
        for key in state_dict.keys():
            if "residual_layer" in key:
                temp_res_num = int(key.split(".")[1])
                if temp_res_num > res_num:
                    res_num = temp_res_num
        res_num = res_num + 1  # zero-indexed

        residual_layer = []
        self.res_num = res_num

        if (
            "residual_layer.0.residual_layer.0.layer.2.fn.rel_pos_bias.weight"
            in state_dict.keys()
        ):
            rel_pos_bias_weight = state_dict[
                "residual_layer.0.residual_layer.0.layer.2.fn.rel_pos_bias.weight"
            ].shape[0]
            self.window_size = int((math.sqrt(rel_pos_bias_weight) + 1) / 2)
        else:
            self.window_size = 8

        self.up_scale = up_scale

        for _ in range(res_num):
            temp_res = OSAG(
                channel_num=num_feat,
                bias=bias,
                block_num=block_num,
                ffn_bias=ffn_bias,
                window_size=self.window_size,
                pe=pe,
            )
            residual_layer.append(temp_res)
        self.residual_layer = nn.Sequential(*residual_layer)
        self.input = nn.Conv2d(
            in_channels=num_in_ch,
            out_channels=num_feat,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=bias,
        )
        self.output = nn.Conv2d(
            in_channels=num_feat,
            out_channels=num_feat,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=bias,
        )
        self.up = pixelshuffle_block(num_feat, num_out_ch, up_scale, bias=bias)

        # self.tail   = pixelshuffle_block(num_feat,num_out_ch,up_scale,bias=bias)

        # for m in self.modules():
        #     if isinstance(m, nn.Conv2d):
        #         n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        #         m.weight.data.normal_(0, sqrt(2. / n))

        # chaiNNer specific stuff
        self.model_arch = "OmniSR"
        self.sub_type = "SR"
        self.in_nc = num_in_ch
        self.out_nc = num_out_ch
        self.num_feat = num_feat
        self.scale = up_scale

        self.supports_fp16 = True  # TODO: Test this
        self.supports_bfp16 = True
        self.min_size_restriction = 16

        self.load_state_dict(state_dict, strict=False)

    def check_image_size(self, x):
        _, _, h, w = x.size()
        # import pdb; pdb.set_trace()
        mod_pad_h = (self.window_size - h % self.window_size) % self.window_size
        mod_pad_w = (self.window_size - w % self.window_size) % self.window_size
        # x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect')
        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "constant", 0)
        return x

    def forward(self, x):
        H, W = x.shape[2:]
        x = self.check_image_size(x)

        residual = self.input(x)
        out = self.residual_layer(residual)

        # origin
        out = torch.add(self.output(out), residual)
        out = self.up(out)

        out = out[:, :, : H * self.up_scale, : W * self.up_scale]
        return out

nodes/impl/pytorch/architecture/OmniSR/OmniSR.py
vladmandic/sd-extension-chainner
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: esa.py
# Created Date: Tuesday April 28th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified:  Thursday, 20th April 2023 9:28:06 am
# Modified By: Chen Xuanhong
# Copyright (c) 2020 Shanghai Jiao Tong University
#############################################################

import torch
import torch.nn as nn
import torch.nn.functional as F

from .layernorm import LayerNorm2d


def moment(x, dim=(2, 3), k=2):
    assert len(x.size()) == 4
    mean = torch.mean(x, dim=dim).unsqueeze(-1).unsqueeze(-1)
    mk = (1 / (x.size(2) * x.size(3))) * torch.sum(torch.pow(x - mean, k), dim=dim)
    return mk


class ESA(nn.Module):
    """
    Modification of Enhanced Spatial Attention (ESA), which is proposed by
    `Residual Feature Aggregation Network for Image Super-Resolution`
    Note: `conv_max` and `conv3_` are NOT used here, so the corresponding codes
    are deleted.
    """

    def __init__(self, esa_channels, n_feats, conv=nn.Conv2d):
        super(ESA, self).__init__()
        f = esa_channels
        self.conv1 = conv(n_feats, f, kernel_size=1)
        self.conv_f = conv(f, f, kernel_size=1)
        self.conv2 = conv(f, f, kernel_size=3, stride=2, padding=0)
        self.conv3 = conv(f, f, kernel_size=3, padding=1)
        self.conv4 = conv(f, n_feats, kernel_size=1)
        self.sigmoid = nn.Sigmoid()
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        c1_ = self.conv1(x)
        c1 = self.conv2(c1_)
        v_max = F.max_pool2d(c1, kernel_size=7, stride=3)
        c3 = self.conv3(v_max)
        c3 = F.interpolate(
            c3, (x.size(2), x.size(3)), mode="bilinear", align_corners=False
        )
        cf = self.conv_f(c1_)
        c4 = self.conv4(c3 + cf)
        m = self.sigmoid(c4)
        return x * m


class LK_ESA(nn.Module):
    def __init__(
        self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
    ):
        super(LK_ESA, self).__init__()
        f = esa_channels
        self.conv1 = conv(n_feats, f, kernel_size=1)
        self.conv_f = conv(f, f, kernel_size=1)

        kernel_size = 17
        kernel_expand = kernel_expand
        padding = kernel_size // 2

        self.vec_conv = nn.Conv2d(
            in_channels=f * kernel_expand,
            out_channels=f * kernel_expand,
            kernel_size=(1, kernel_size),
            padding=(0, padding),
            groups=2,
            bias=bias,
        )
        self.vec_conv3x1 = nn.Conv2d(
            in_channels=f * kernel_expand,
            out_channels=f * kernel_expand,
            kernel_size=(1, 3),
            padding=(0, 1),
            groups=2,
            bias=bias,
        )

        self.hor_conv = nn.Conv2d(
            in_channels=f * kernel_expand,
            out_channels=f * kernel_expand,
            kernel_size=(kernel_size, 1),
            padding=(padding, 0),
            groups=2,
            bias=bias,
        )
        self.hor_conv1x3 = nn.Conv2d(
            in_channels=f * kernel_expand,
            out_channels=f * kernel_expand,
            kernel_size=(3, 1),
            padding=(1, 0),
            groups=2,
            bias=bias,
        )

        self.conv4 = conv(f, n_feats, kernel_size=1)
        self.sigmoid = nn.Sigmoid()
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        c1_ = self.conv1(x)

        res = self.vec_conv(c1_) + self.vec_conv3x1(c1_)
        res = self.hor_conv(res) + self.hor_conv1x3(res)

        cf = self.conv_f(c1_)
        c4 = self.conv4(res + cf)
        m = self.sigmoid(c4)
        return x * m


class LK_ESA_LN(nn.Module):
    def __init__(
        self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
    ):
        super(LK_ESA_LN, self).__init__()
        f = esa_channels
        self.conv1 = conv(n_feats, f, kernel_size=1)
        self.conv_f = conv(f, f, kernel_size=1)

        kernel_size = 17
        kernel_expand = kernel_expand
        padding = kernel_size // 2

        self.norm = LayerNorm2d(n_feats)

        self.vec_conv = nn.Conv2d(
            in_channels=f * kernel_expand,
            out_channels=f * kernel_expand,
            kernel_size=(1, kernel_size),
            padding=(0, padding),
            groups=2,
            bias=bias,
        )
        self.vec_conv3x1 = nn.Conv2d(
            in_channels=f * kernel_expand,
            out_channels=f * kernel_expand,
            kernel_size=(1, 3),
            padding=(0, 1),
            groups=2,
            bias=bias,
        )

        self.hor_conv = nn.Conv2d(
            in_channels=f * kernel_expand,
            out_channels=f * kernel_expand,
            kernel_size=(kernel_size, 1),
            padding=(padding, 0),
            groups=2,
            bias=bias,
        )
        self.hor_conv1x3 = nn.Conv2d(
            in_channels=f * kernel_expand,
            out_channels=f * kernel_expand,
            kernel_size=(3, 1),
            padding=(1, 0),
            groups=2,
            bias=bias,
        )

        self.conv4 = conv(f, n_feats, kernel_size=1)
        self.sigmoid = nn.Sigmoid()
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        c1_ = self.norm(x)
        c1_ = self.conv1(c1_)

        res = self.vec_conv(c1_) + self.vec_conv3x1(c1_)
        res = self.hor_conv(res) + self.hor_conv1x3(res)

        cf = self.conv_f(c1_)
        c4 = self.conv4(res + cf)
        m = self.sigmoid(c4)
        return x * m


class AdaGuidedFilter(nn.Module):
    def __init__(
        self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
    ):
        super(AdaGuidedFilter, self).__init__()

        self.gap = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Conv2d(
            in_channels=n_feats,
            out_channels=1,
            kernel_size=1,
            padding=0,
            stride=1,
            groups=1,
            bias=True,
        )

        self.r = 5

    def box_filter(self, x, r):
        channel = x.shape[1]
        kernel_size = 2 * r + 1
        weight = 1.0 / (kernel_size**2)
        box_kernel = weight * torch.ones(
            (channel, 1, kernel_size, kernel_size), dtype=torch.float32, device=x.device
        )
        output = F.conv2d(x, weight=box_kernel, stride=1, padding=r, groups=channel)
        return output

    def forward(self, x):
        _, _, H, W = x.shape
        N = self.box_filter(
            torch.ones((1, 1, H, W), dtype=x.dtype, device=x.device), self.r
        )

        # epsilon = self.fc(self.gap(x))
        # epsilon = torch.pow(epsilon, 2)
        epsilon = 1e-2

        mean_x = self.box_filter(x, self.r) / N
        var_x = self.box_filter(x * x, self.r) / N - mean_x * mean_x

        A = var_x / (var_x + epsilon)
        b = (1 - A) * mean_x
        m = A * x + b

        # mean_A = self.box_filter(A, self.r) / N
        # mean_b = self.box_filter(b, self.r) / N
        # m = mean_A * x + mean_b
        return x * m


class AdaConvGuidedFilter(nn.Module):
    def __init__(
        self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
    ):
        super(AdaConvGuidedFilter, self).__init__()
        f = esa_channels

        self.conv_f = conv(f, f, kernel_size=1)

        kernel_size = 17
        kernel_expand = kernel_expand
        padding = kernel_size // 2

        self.vec_conv = nn.Conv2d(
            in_channels=f,
            out_channels=f,
            kernel_size=(1, kernel_size),
            padding=(0, padding),
            groups=f,
            bias=bias,
        )

        self.hor_conv = nn.Conv2d(
            in_channels=f,
            out_channels=f,
            kernel_size=(kernel_size, 1),
            padding=(padding, 0),
            groups=f,
            bias=bias,
        )

        self.gap = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Conv2d(
            in_channels=f,
            out_channels=f,
            kernel_size=1,
            padding=0,
            stride=1,
            groups=1,
            bias=True,
        )

    def forward(self, x):
        y = self.vec_conv(x)
        y = self.hor_conv(y)

        sigma = torch.pow(y, 2)
        epsilon = self.fc(self.gap(y))

        weight = sigma / (sigma + epsilon)

        m = weight * x + (1 - weight)

        return x * m

nodes/impl/pytorch/architecture/OmniSR/esa.py
vladmandic/sd-extension-chainner
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: layernorm.py
# Created Date: Tuesday April 28th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified:  Thursday, 20th April 2023 9:28:20 am
# Modified By: Chen Xuanhong
# Copyright (c) 2020 Shanghai Jiao Tong University
#############################################################

import torch
import torch.nn as nn


class LayerNormFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias, eps):
        ctx.eps = eps
        N, C, H, W = x.size()
        mu = x.mean(1, keepdim=True)
        var = (x - mu).pow(2).mean(1, keepdim=True)
        y = (x - mu) / (var + eps).sqrt()
        ctx.save_for_backward(y, var, weight)
        y = weight.view(1, C, 1, 1) * y + bias.view(1, C, 1, 1)
        return y

    @staticmethod
    def backward(ctx, grad_output):
        eps = ctx.eps

        N, C, H, W = grad_output.size()
        y, var, weight = ctx.saved_variables
        g = grad_output * weight.view(1, C, 1, 1)
        mean_g = g.mean(dim=1, keepdim=True)

        mean_gy = (g * y).mean(dim=1, keepdim=True)
        gx = 1.0 / torch.sqrt(var + eps) * (g - y * mean_gy - mean_g)
        return (
            gx,
            (grad_output * y).sum(dim=3).sum(dim=2).sum(dim=0),
            grad_output.sum(dim=3).sum(dim=2).sum(dim=0),
            None,
        )


class LayerNorm2d(nn.Module):
    def __init__(self, channels, eps=1e-6):
        super(LayerNorm2d, self).__init__()
        self.register_parameter("weight", nn.Parameter(torch.ones(channels)))
        self.register_parameter("bias", nn.Parameter(torch.zeros(channels)))
        self.eps = eps

    def forward(self, x):
        return LayerNormFunction.apply(x, self.weight, self.bias, self.eps)


class GRN(nn.Module):
    """GRN (Global Response Normalization) layer"""

    def __init__(self, dim):
        super().__init__()
        self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1))
        self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1))

    def forward(self, x):
        Gx = torch.norm(x, p=2, dim=(2, 3), keepdim=True)
        Nx = Gx / (Gx.mean(dim=1, keepdim=True) + 1e-6)
        return self.gamma * (x * Nx) + self.beta + x

nodes/impl/pytorch/architecture/OmniSR/layernorm.py
vladmandic/sd-extension-chainner
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: pixelshuffle.py
# Created Date: Friday July 1st 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified:  Friday, 1st July 2022 10:18:39 am
# Modified By: Chen Xuanhong
# Copyright (c) 2022 Shanghai Jiao Tong University
#############################################################

import torch.nn as nn


def pixelshuffle_block(
    in_channels, out_channels, upscale_factor=2, kernel_size=3, bias=False
):
    """
    Upsample features according to `upscale_factor`.
    """
    kernel_size // 2
    conv = nn.Conv2d(
        in_channels,
        out_channels * (upscale_factor**2),
        kernel_size,
        padding=1,
        bias=bias,
    )
    pixel_shuffle = nn.PixelShuffle(upscale_factor)
    return nn.Sequential(*[conv, pixel_shuffle])

nodes/impl/pytorch/architecture/OmniSR/pixelshuffle.py
vladmandic/sd-extension-chainner
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import functools
import math
import re
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.nn.functional as F

from . import block as B


# Borrowed from https://github.com/rlaphoenix/VSGAN/blob/master/vsgan/archs/ESRGAN.py
# Which enhanced stuff that was already here
class RRDBNet(nn.Module):
    def __init__(
        self,
        state_dict,
        norm=None,
        act: str = "leakyrelu",
        upsampler: str = "upconv",
        mode: B.ConvMode = "CNA",
    ) -> None:
        """
        ESRGAN - Enhanced Super-Resolution Generative Adversarial Networks.
        By Xintao Wang, Ke Yu, Shixiang Wu, Jinjin Gu, Yihao Liu, Chao Dong, Yu Qiao,
        and Chen Change Loy.
        This is old-arch Residual in Residual Dense Block Network and is not
        the newest revision that's available at github.com/xinntao/ESRGAN.
        This is on purpose, the newest Network has severely limited the
        potential use of the Network with no benefits.
        This network supports model files from both new and old-arch.
        Args:
            norm: Normalization layer
            act: Activation layer
            upsampler: Upsample layer. upconv, pixel_shuffle
            mode: Convolution mode
        """
        super(RRDBNet, self).__init__()
        self.model_arch = "ESRGAN"
        self.sub_type = "SR"

        self.state = state_dict
        self.norm = norm
        self.act = act
        self.upsampler = upsampler
        self.mode = mode

        self.state_map = {
            # currently supports old, new, and newer RRDBNet arch models
            # ESRGAN, BSRGAN/RealSR, Real-ESRGAN
            "model.0.weight": ("conv_first.weight",),
            "model.0.bias": ("conv_first.bias",),
            "model.1.sub./NB/.weight": ("trunk_conv.weight", "conv_body.weight"),
            "model.1.sub./NB/.bias": ("trunk_conv.bias", "conv_body.bias"),
            r"model.1.sub.\1.RDB\2.conv\3.0.\4": (
                r"RRDB_trunk\.(\d+)\.RDB(\d)\.conv(\d+)\.(weight|bias)",
                r"body\.(\d+)\.rdb(\d)\.conv(\d+)\.(weight|bias)",
            ),
        }
        if "params_ema" in self.state:
            self.state = self.state["params_ema"]
            # self.model_arch = "RealESRGAN"
        self.num_blocks = self.get_num_blocks()
        self.plus = any("conv1x1" in k for k in self.state.keys())
        if self.plus:
            self.model_arch = "ESRGAN+"

        self.state = self.new_to_old_arch(self.state)

        self.key_arr = list(self.state.keys())

        self.in_nc: int = self.state[self.key_arr[0]].shape[1]
        self.out_nc: int = self.state[self.key_arr[-1]].shape[0]

        self.scale: int = self.get_scale()
        self.num_filters: int = self.state[self.key_arr[0]].shape[0]

        c2x2 = False
        if self.state["model.0.weight"].shape[-2] == 2:
            c2x2 = True
            self.scale = round(math.sqrt(self.scale / 4))
            self.model_arch = "ESRGAN-2c2"

        self.supports_fp16 = True
        self.supports_bfp16 = True
        self.min_size_restriction = None

        # Detect if pixelunshuffle was used (Real-ESRGAN)
        if self.in_nc in (self.out_nc * 4, self.out_nc * 16) and self.out_nc in (
            self.in_nc / 4,
            self.in_nc / 16,
        ):
            self.shuffle_factor = int(math.sqrt(self.in_nc / self.out_nc))
        else:
            self.shuffle_factor = None

        upsample_block = {
            "upconv": B.upconv_block,
            "pixel_shuffle": B.pixelshuffle_block,
        }.get(self.upsampler)
        if upsample_block is None:
            raise NotImplementedError(f"Upsample mode [{self.upsampler}] is not found")

        if self.scale == 3:
            upsample_blocks = upsample_block(
                in_nc=self.num_filters,
                out_nc=self.num_filters,
                upscale_factor=3,
                act_type=self.act,
                c2x2=c2x2,
            )
        else:
            upsample_blocks = [
                upsample_block(
                    in_nc=self.num_filters,
                    out_nc=self.num_filters,
                    act_type=self.act,
                    c2x2=c2x2,
                )
                for _ in range(int(math.log(self.scale, 2)))
            ]

        self.model = B.sequential(
            # fea conv
            B.conv_block(
                in_nc=self.in_nc,
                out_nc=self.num_filters,
                kernel_size=3,
                norm_type=None,
                act_type=None,
                c2x2=c2x2,
            ),
            B.ShortcutBlock(
                B.sequential(
                    # rrdb blocks
                    *[
                        B.RRDB(
                            nf=self.num_filters,
                            kernel_size=3,
                            gc=32,
                            stride=1,
                            bias=True,
                            pad_type="zero",
                            norm_type=self.norm,
                            act_type=self.act,
                            mode="CNA",
                            plus=self.plus,
                            c2x2=c2x2,
                        )
                        for _ in range(self.num_blocks)
                    ],
                    # lr conv
                    B.conv_block(
                        in_nc=self.num_filters,
                        out_nc=self.num_filters,
                        kernel_size=3,
                        norm_type=self.norm,
                        act_type=None,
                        mode=self.mode,
                        c2x2=c2x2,
                    ),
                )
            ),
            *upsample_blocks,
            # hr_conv0
            B.conv_block(
                in_nc=self.num_filters,
                out_nc=self.num_filters,
                kernel_size=3,
                norm_type=None,
                act_type=self.act,
                c2x2=c2x2,
            ),
            # hr_conv1
            B.conv_block(
                in_nc=self.num_filters,
                out_nc=self.out_nc,
                kernel_size=3,
                norm_type=None,
                act_type=None,
                c2x2=c2x2,
            ),
        )

        # Adjust these properties for calculations outside of the model
        if self.shuffle_factor:
            self.in_nc //= self.shuffle_factor**2
            self.scale //= self.shuffle_factor

        self.load_state_dict(self.state, strict=False)

    def new_to_old_arch(self, state):
        """Convert a new-arch model state dictionary to an old-arch dictionary."""
        if "params_ema" in state:
            state = state["params_ema"]

        if "conv_first.weight" not in state:
            # model is already old arch, this is a loose check, but should be sufficient
            return state

        # add nb to state keys
        for kind in ("weight", "bias"):
            self.state_map[f"model.1.sub.{self.num_blocks}.{kind}"] = self.state_map[
                f"model.1.sub./NB/.{kind}"
            ]
            del self.state_map[f"model.1.sub./NB/.{kind}"]

        old_state = OrderedDict()
        for old_key, new_keys in self.state_map.items():
            for new_key in new_keys:
                if r"\1" in old_key:
                    for k, v in state.items():
                        sub = re.sub(new_key, old_key, k)
                        if sub != k:
                            old_state[sub] = v
                else:
                    if new_key in state:
                        old_state[old_key] = state[new_key]

        # upconv layers
        max_upconv = 0
        for key in state.keys():
            match = re.match(r"(upconv|conv_up)(\d)\.(weight|bias)", key)
            if match is not None:
                _, key_num, key_type = match.groups()
                old_state[f"model.{int(key_num) * 3}.{key_type}"] = state[key]
                max_upconv = max(max_upconv, int(key_num) * 3)

        # final layers
        for key in state.keys():
            if key in ("HRconv.weight", "conv_hr.weight"):
                old_state[f"model.{max_upconv + 2}.weight"] = state[key]
            elif key in ("HRconv.bias", "conv_hr.bias"):
                old_state[f"model.{max_upconv + 2}.bias"] = state[key]
            elif key in ("conv_last.weight",):
                old_state[f"model.{max_upconv + 4}.weight"] = state[key]
            elif key in ("conv_last.bias",):
                old_state[f"model.{max_upconv + 4}.bias"] = state[key]

        # Sort by first numeric value of each layer
        def compare(item1, item2):
            parts1 = item1.split(".")
            parts2 = item2.split(".")
            int1 = int(parts1[1])
            int2 = int(parts2[1])
            return int1 - int2

        sorted_keys = sorted(old_state.keys(), key=functools.cmp_to_key(compare))

        # Rebuild the output dict in the right order
        out_dict = OrderedDict((k, old_state[k]) for k in sorted_keys)

        return out_dict

    def get_scale(self, min_part: int = 6) -> int:
        n = 0
        for part in list(self.state):
            parts = part.split(".")[1:]
            if len(parts) == 2:
                part_num = int(parts[0])
                if part_num > min_part and parts[1] == "weight":
                    n += 1
        return 2**n

    def get_num_blocks(self) -> int:
        nbs = []
        state_keys = self.state_map[r"model.1.sub.\1.RDB\2.conv\3.0.\4"] + (
            r"model\.\d+\.sub\.(\d+)\.RDB(\d+)\.conv(\d+)\.0\.(weight|bias)",
        )
        for state_key in state_keys:
            for k in self.state:
                m = re.search(state_key, k)
                if m:
                    nbs.append(int(m.group(1)))
            if nbs:
                break
        return max(*nbs) + 1

    def forward(self, x):
        if self.shuffle_factor:
            _, _, h, w = x.size()
            mod_pad_h = (
                self.shuffle_factor - h % self.shuffle_factor
            ) % self.shuffle_factor
            mod_pad_w = (
                self.shuffle_factor - w % self.shuffle_factor
            ) % self.shuffle_factor
            x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect")
            x = torch.pixel_unshuffle(x, downscale_factor=self.shuffle_factor)
            x = self.model(x)
            return x[:, :, : h * self.scale, : w * self.scale]
        return self.model(x)

nodes/impl/pytorch/architecture/RRDB.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# -----------------------------------------------------------------------------------
# SCUNet: Practical Blind Denoising via Swin-Conv-UNet and Data Synthesis, https://arxiv.org/abs/2203.13278
# Zhang, Kai and Li, Yawei and Liang, Jingyun and Cao, Jiezhang and Zhang, Yulun and Tang, Hao and Timofte, Radu and Van Gool, Luc
# -----------------------------------------------------------------------------------

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from einops.layers.torch import Rearrange

from nodes.utils.utils import round_half_up

from .timm.drop import DropPath
from .timm.weight_init import trunc_normal_


# Borrowed from https://github.com/cszn/SCUNet/blob/main/models/network_scunet.py
class WMSA(nn.Module):
    """Self-attention module in Swin Transformer"""

    def __init__(self, input_dim, output_dim, head_dim, window_size, type):
        super(WMSA, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.head_dim = head_dim
        self.scale = self.head_dim**-0.5
        self.n_heads = input_dim // head_dim
        self.window_size = window_size
        self.type = type
        self.embedding_layer = nn.Linear(self.input_dim, 3 * self.input_dim, bias=True)

        self.relative_position_params = nn.Parameter(
            torch.zeros((2 * window_size - 1) * (2 * window_size - 1), self.n_heads)
        )
        # TODO recover
        # self.relative_position_params = nn.Parameter(torch.zeros(self.n_heads, 2 * window_size - 1, 2 * window_size -1))
        self.relative_position_params = nn.Parameter(
            torch.zeros((2 * window_size - 1) * (2 * window_size - 1), self.n_heads)
        )

        self.linear = nn.Linear(self.input_dim, self.output_dim)

        trunc_normal_(self.relative_position_params, std=0.02)
        self.relative_position_params = torch.nn.Parameter(
            self.relative_position_params.view(
                2 * window_size - 1, 2 * window_size - 1, self.n_heads
            )
            .transpose(1, 2)
            .transpose(0, 1)
        )

    def generate_mask(self, h, w, p, shift):
        """generating the mask of SW-MSA
        Args:
            shift: shift parameters in CyclicShift.
        Returns:
            attn_mask: should be (1 1 w p p),
        """
        # supporting square.
        attn_mask = torch.zeros(
            h,
            w,
            p,
            p,
            p,
            p,
            dtype=torch.bool,
            device=self.relative_position_params.device,
        )
        if self.type == "W":
            return attn_mask

        s = p - shift
        attn_mask[-1, :, :s, :, s:, :] = True
        attn_mask[-1, :, s:, :, :s, :] = True
        attn_mask[:, -1, :, :s, :, s:] = True
        attn_mask[:, -1, :, s:, :, :s] = True
        attn_mask = rearrange(
            attn_mask, "w1 w2 p1 p2 p3 p4 -> 1 1 (w1 w2) (p1 p2) (p3 p4)"
        )
        return attn_mask

    def forward(self, x):
        """Forward pass of Window Multi-head Self-attention module.
        Args:
            x: input tensor with shape of [b h w c];
            attn_mask: attention mask, fill -inf where the value is True;
        Returns:
            output: tensor shape [b h w c]
        """
        if self.type != "W":
            x = torch.roll(
                x,
                shifts=(-(self.window_size // 2), -(self.window_size // 2)),
                dims=(1, 2),
            )

        x = rearrange(
            x,
            "b (w1 p1) (w2 p2) c -> b w1 w2 p1 p2 c",
            p1=self.window_size,
            p2=self.window_size,
        )
        h_windows = x.size(1)
        w_windows = x.size(2)
        # square validation
        # assert h_windows == w_windows

        x = rearrange(
            x,
            "b w1 w2 p1 p2 c -> b (w1 w2) (p1 p2) c",
            p1=self.window_size,
            p2=self.window_size,
        )
        qkv = self.embedding_layer(x)
        q, k, v = rearrange(
            qkv, "b nw np (threeh c) -> threeh b nw np c", c=self.head_dim
        ).chunk(3, dim=0)
        sim = torch.einsum("hbwpc,hbwqc->hbwpq", q, k) * self.scale
        # Adding learnable relative embedding
        sim = sim + rearrange(self.relative_embedding(), "h p q -> h 1 1 p q")
        # Using Attn Mask to distinguish different subwindows.
        if self.type != "W":
            attn_mask = self.generate_mask(
                h_windows, w_windows, self.window_size, shift=self.window_size // 2
            )
            sim = sim.masked_fill_(attn_mask, float("-inf"))

        probs = nn.functional.softmax(sim, dim=-1)
        output = torch.einsum("hbwij,hbwjc->hbwic", probs, v)
        output = rearrange(output, "h b w p c -> b w p (h c)")
        output = self.linear(output)
        output = rearrange(
            output,
            "b (w1 w2) (p1 p2) c -> b (w1 p1) (w2 p2) c",
            w1=h_windows,
            p1=self.window_size,
        )

        if self.type != "W":
            output = torch.roll(
                output,
                shifts=(self.window_size // 2, self.window_size // 2),
                dims=(1, 2),
            )

        return output

    def relative_embedding(self):
        cord = torch.tensor(
            np.array(
                [
                    [i, j]
                    for i in range(self.window_size)
                    for j in range(self.window_size)
                ]
            )
        )
        relation = cord[:, None, :] - cord[None, :, :] + self.window_size - 1
        # negative is allowed
        return self.relative_position_params[
            :, relation[:, :, 0].long(), relation[:, :, 1].long()
        ]


class Block(nn.Module):
    def __init__(
        self,
        input_dim,
        output_dim,
        head_dim,
        window_size,
        drop_path,
        type="W",
        input_resolution=None,
    ):
        """SwinTransformer Block"""
        super(Block, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        assert type in ["W", "SW"]
        self.type = type
        if input_resolution <= window_size:
            self.type = "W"

        self.ln1 = nn.LayerNorm(input_dim)
        self.msa = WMSA(input_dim, input_dim, head_dim, window_size, self.type)
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.ln2 = nn.LayerNorm(input_dim)
        self.mlp = nn.Sequential(
            nn.Linear(input_dim, 4 * input_dim),
            nn.GELU(),
            nn.Linear(4 * input_dim, output_dim),
        )

    def forward(self, x):
        x = x + self.drop_path(self.msa(self.ln1(x)))
        x = x + self.drop_path(self.mlp(self.ln2(x)))
        return x


class ConvTransBlock(nn.Module):
    def __init__(
        self,
        conv_dim,
        trans_dim,
        head_dim,
        window_size,
        drop_path,
        type="W",
        input_resolution=None,
    ):
        """SwinTransformer and Conv Block"""
        super(ConvTransBlock, self).__init__()
        self.conv_dim = conv_dim
        self.trans_dim = trans_dim
        self.head_dim = head_dim
        self.window_size = window_size
        self.drop_path = drop_path
        self.type = type
        self.input_resolution = input_resolution

        assert self.type in ["W", "SW"]
        if self.input_resolution <= self.window_size:
            self.type = "W"

        self.trans_block = Block(
            self.trans_dim,
            self.trans_dim,
            self.head_dim,
            self.window_size,
            self.drop_path,
            self.type,
            self.input_resolution,
        )
        self.conv1_1 = nn.Conv2d(
            self.conv_dim + self.trans_dim,
            self.conv_dim + self.trans_dim,
            1,
            1,
            0,
            bias=True,
        )
        self.conv1_2 = nn.Conv2d(
            self.conv_dim + self.trans_dim,
            self.conv_dim + self.trans_dim,
            1,
            1,
            0,
            bias=True,
        )

        self.conv_block = nn.Sequential(
            nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False),
            nn.ReLU(True),
            nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False),
        )

    def forward(self, x):
        conv_x, trans_x = torch.split(
            self.conv1_1(x), (self.conv_dim, self.trans_dim), dim=1
        )
        conv_x = self.conv_block(conv_x) + conv_x
        trans_x = Rearrange("b c h w -> b h w c")(trans_x)
        trans_x = self.trans_block(trans_x)
        trans_x = Rearrange("b h w c -> b c h w")(trans_x)
        res = self.conv1_2(torch.cat((conv_x, trans_x), dim=1))
        x = x + res

        return x


class SCUNet(nn.Module):
    def __init__(
        self,
        state_dict,
        in_nc=3,
        config=None,
        dim=64,
        drop_path_rate=0.0,
        input_resolution=256,
    ):
        if config is None:
            config = [4, 4, 4, 4, 4, 4, 4]
        super(SCUNet, self).__init__()
        self.model_arch = "SCUNet"
        self.sub_type = "SR"

        self.num_filters: int = 0

        self.state = state_dict
        self.config = config
        self.dim = dim
        self.head_dim = 32
        self.window_size = 8

        self.in_nc = in_nc
        self.out_nc = self.in_nc
        self.scale = 1
        self.supports_fp16 = True

        # drop path rate for each layer
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(config))]

        self.m_head = [nn.Conv2d(in_nc, dim, 3, 1, 1, bias=False)]

        begin = 0
        self.m_down1 = [
            ConvTransBlock(
                dim // 2,
                dim // 2,
                self.head_dim,
                self.window_size,
                dpr[i + begin],
                "W" if not i % 2 else "SW",
                input_resolution,
            )
            for i in range(config[0])
        ] + [nn.Conv2d(dim, 2 * dim, 2, 2, 0, bias=False)]

        begin += config[0]
        self.m_down2 = [
            ConvTransBlock(
                dim,
                dim,
                self.head_dim,
                self.window_size,
                dpr[i + begin],
                "W" if not i % 2 else "SW",
                input_resolution // 2,
            )
            for i in range(config[1])
        ] + [nn.Conv2d(2 * dim, 4 * dim, 2, 2, 0, bias=False)]

        begin += config[1]
        self.m_down3 = [
            ConvTransBlock(
                2 * dim,
                2 * dim,
                self.head_dim,
                self.window_size,
                dpr[i + begin],
                "W" if not i % 2 else "SW",
                input_resolution // 4,
            )
            for i in range(config[2])
        ] + [nn.Conv2d(4 * dim, 8 * dim, 2, 2, 0, bias=False)]

        begin += config[2]
        self.m_body = [
            ConvTransBlock(
                4 * dim,
                4 * dim,
                self.head_dim,
                self.window_size,
                dpr[i + begin],
                "W" if not i % 2 else "SW",
                input_resolution // 8,
            )
            for i in range(config[3])
        ]

        begin += config[3]
        self.m_up3 = [
            nn.ConvTranspose2d(8 * dim, 4 * dim, 2, 2, 0, bias=False),
        ] + [
            ConvTransBlock(
                2 * dim,
                2 * dim,
                self.head_dim,
                self.window_size,
                dpr[i + begin],
                "W" if not i % 2 else "SW",
                input_resolution // 4,
            )
            for i in range(config[4])
        ]

        begin += config[4]
        self.m_up2 = [
            nn.ConvTranspose2d(4 * dim, 2 * dim, 2, 2, 0, bias=False),
        ] + [
            ConvTransBlock(
                dim,
                dim,
                self.head_dim,
                self.window_size,
                dpr[i + begin],
                "W" if not i % 2 else "SW",
                input_resolution // 2,
            )
            for i in range(config[5])
        ]

        begin += config[5]
        self.m_up1 = [
            nn.ConvTranspose2d(2 * dim, dim, 2, 2, 0, bias=False),
        ] + [
            ConvTransBlock(
                dim // 2,
                dim // 2,
                self.head_dim,
                self.window_size,
                dpr[i + begin],
                "W" if not i % 2 else "SW",
                input_resolution,
            )
            for i in range(config[6])
        ]

        self.m_tail = [nn.Conv2d(dim, in_nc, 3, 1, 1, bias=False)]

        self.m_head = nn.Sequential(*self.m_head)
        self.m_down1 = nn.Sequential(*self.m_down1)
        self.m_down2 = nn.Sequential(*self.m_down2)
        self.m_down3 = nn.Sequential(*self.m_down3)
        self.m_body = nn.Sequential(*self.m_body)
        self.m_up3 = nn.Sequential(*self.m_up3)
        self.m_up2 = nn.Sequential(*self.m_up2)
        self.m_up1 = nn.Sequential(*self.m_up1)
        self.m_tail = nn.Sequential(*self.m_tail)
        # self.apply(self._init_weights)
        self.load_state_dict(state_dict, strict=True)

    def check_image_size(self, x):
        _, _, h, w = x.size()
        mod_pad_h = (64 - h % 64) % 64
        mod_pad_w = (64 - w % 64) % 64
        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect")
        return x

    def forward(self, x0):
        h, w = x0.size()[-2:]
        x0 = self.check_image_size(x0)

        x1 = self.m_head(x0)
        x2 = self.m_down1(x1)
        x3 = self.m_down2(x2)
        x4 = self.m_down3(x3)
        x = self.m_body(x4)
        x = self.m_up3(x + x4)
        x = self.m_up2(x + x3)
        x = self.m_up1(x + x2)
        x = self.m_tail(x + x1)

        x = x[:, :, :h, :w]
        return x

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

nodes/impl/pytorch/architecture/SCUNet.py
vladmandic/sd-extension-chainner
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from . import block as B


class Get_gradient_nopadding(nn.Module):
    def __init__(self):
        super(Get_gradient_nopadding, self).__init__()
        kernel_v = [[0, -1, 0], [0, 0, 0], [0, 1, 0]]
        kernel_h = [[0, 0, 0], [-1, 0, 1], [0, 0, 0]]
        kernel_h = torch.FloatTensor(kernel_h).unsqueeze(0).unsqueeze(0)
        kernel_v = torch.FloatTensor(kernel_v).unsqueeze(0).unsqueeze(0)
        self.weight_h = nn.Parameter(data=kernel_h, requires_grad=False)  # type: ignore

        self.weight_v = nn.Parameter(data=kernel_v, requires_grad=False)  # type: ignore

    def forward(self, x):
        x_list = []
        for i in range(x.shape[1]):
            x_i = x[:, i]
            x_i_v = F.conv2d(x_i.unsqueeze(1), self.weight_v, padding=1)
            x_i_h = F.conv2d(x_i.unsqueeze(1), self.weight_h, padding=1)
            x_i = torch.sqrt(torch.pow(x_i_v, 2) + torch.pow(x_i_h, 2) + 1e-6)
            x_list.append(x_i)

        x = torch.cat(x_list, dim=1)

        return x


class SPSRNet(nn.Module):
    def __init__(
        self,
        state_dict,
        norm=None,
        act: str = "leakyrelu",
        upsampler: str = "upconv",
        mode: B.ConvMode = "CNA",
    ):
        super(SPSRNet, self).__init__()
        self.model_arch = "SPSR"
        self.sub_type = "SR"

        self.state = state_dict
        self.norm = norm
        self.act = act
        self.upsampler = upsampler
        self.mode = mode

        self.num_blocks = self.get_num_blocks()

        self.in_nc: int = self.state["model.0.weight"].shape[1]
        self.out_nc: int = self.state["f_HR_conv1.0.bias"].shape[0]

        self.scale = self.get_scale(4)
        self.num_filters: int = self.state["model.0.weight"].shape[0]

        self.supports_fp16 = True
        self.supports_bfp16 = True
        self.min_size_restriction = None

        n_upscale = int(math.log(self.scale, 2))
        if self.scale == 3:
            n_upscale = 1

        fea_conv = B.conv_block(
            self.in_nc, self.num_filters, kernel_size=3, norm_type=None, act_type=None
        )
        rb_blocks = [
            B.RRDB(
                self.num_filters,
                kernel_size=3,
                gc=32,
                stride=1,
                bias=True,
                pad_type="zero",
                norm_type=norm,
                act_type=act,
                mode="CNA",
            )
            for _ in range(self.num_blocks)
        ]
        LR_conv = B.conv_block(
            self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=norm,
            act_type=None,
            mode=mode,
        )

        if upsampler == "upconv":
            upsample_block = B.upconv_block
        elif upsampler == "pixelshuffle":
            upsample_block = B.pixelshuffle_block
        else:
            raise NotImplementedError(f"upsample mode [{upsampler}] is not found")
        if self.scale == 3:
            a_upsampler = upsample_block(
                self.num_filters, self.num_filters, 3, act_type=act
            )
        else:
            a_upsampler = [
                upsample_block(self.num_filters, self.num_filters, act_type=act)
                for _ in range(n_upscale)
            ]
        self.HR_conv0_new = B.conv_block(
            self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=None,
            act_type=act,
        )
        self.HR_conv1_new = B.conv_block(
            self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=None,
            act_type=None,
        )

        self.model = B.sequential(
            fea_conv,
            B.ShortcutBlockSPSR(B.sequential(*rb_blocks, LR_conv)),
            *a_upsampler,
            self.HR_conv0_new,
        )

        self.get_g_nopadding = Get_gradient_nopadding()

        self.b_fea_conv = B.conv_block(
            self.in_nc, self.num_filters, kernel_size=3, norm_type=None, act_type=None
        )

        self.b_concat_1 = B.conv_block(
            2 * self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=None,
            act_type=None,
        )
        self.b_block_1 = B.RRDB(
            self.num_filters * 2,
            kernel_size=3,
            gc=32,
            stride=1,
            bias=True,
            pad_type="zero",
            norm_type=norm,
            act_type=act,
            mode="CNA",
        )

        self.b_concat_2 = B.conv_block(
            2 * self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=None,
            act_type=None,
        )
        self.b_block_2 = B.RRDB(
            self.num_filters * 2,
            kernel_size=3,
            gc=32,
            stride=1,
            bias=True,
            pad_type="zero",
            norm_type=norm,
            act_type=act,
            mode="CNA",
        )

        self.b_concat_3 = B.conv_block(
            2 * self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=None,
            act_type=None,
        )
        self.b_block_3 = B.RRDB(
            self.num_filters * 2,
            kernel_size=3,
            gc=32,
            stride=1,
            bias=True,
            pad_type="zero",
            norm_type=norm,
            act_type=act,
            mode="CNA",
        )

        self.b_concat_4 = B.conv_block(
            2 * self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=None,
            act_type=None,
        )
        self.b_block_4 = B.RRDB(
            self.num_filters * 2,
            kernel_size=3,
            gc=32,
            stride=1,
            bias=True,
            pad_type="zero",
            norm_type=norm,
            act_type=act,
            mode="CNA",
        )

        self.b_LR_conv = B.conv_block(
            self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=norm,
            act_type=None,
            mode=mode,
        )

        if upsampler == "upconv":
            upsample_block = B.upconv_block
        elif upsampler == "pixelshuffle":
            upsample_block = B.pixelshuffle_block
        else:
            raise NotImplementedError(f"upsample mode [{upsampler}] is not found")
        if self.scale == 3:
            b_upsampler = upsample_block(
                self.num_filters, self.num_filters, 3, act_type=act
            )
        else:
            b_upsampler = [
                upsample_block(self.num_filters, self.num_filters, act_type=act)
                for _ in range(n_upscale)
            ]

        b_HR_conv0 = B.conv_block(
            self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=None,
            act_type=act,
        )
        b_HR_conv1 = B.conv_block(
            self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=None,
            act_type=None,
        )

        self.b_module = B.sequential(*b_upsampler, b_HR_conv0, b_HR_conv1)

        self.conv_w = B.conv_block(
            self.num_filters, self.out_nc, kernel_size=1, norm_type=None, act_type=None
        )

        self.f_concat = B.conv_block(
            self.num_filters * 2,
            self.num_filters,
            kernel_size=3,
            norm_type=None,
            act_type=None,
        )

        self.f_block = B.RRDB(
            self.num_filters * 2,
            kernel_size=3,
            gc=32,
            stride=1,
            bias=True,
            pad_type="zero",
            norm_type=norm,
            act_type=act,
            mode="CNA",
        )

        self.f_HR_conv0 = B.conv_block(
            self.num_filters,
            self.num_filters,
            kernel_size=3,
            norm_type=None,
            act_type=act,
        )
        self.f_HR_conv1 = B.conv_block(
            self.num_filters, self.out_nc, kernel_size=3, norm_type=None, act_type=None
        )

        self.load_state_dict(self.state, strict=False)

    def get_scale(self, min_part: int = 4) -> int:
        n = 0
        for part in list(self.state):
            parts = part.split(".")
            if len(parts) == 3:
                part_num = int(parts[1])
                if part_num > min_part and parts[0] == "model" and parts[2] == "weight":
                    n += 1
        return 2**n

    def get_num_blocks(self) -> int:
        nb = 0
        for part in list(self.state):
            parts = part.split(".")
            n_parts = len(parts)
            if n_parts == 5 and parts[2] == "sub":
                nb = int(parts[3])
        return nb

    def forward(self, x):
        x_grad = self.get_g_nopadding(x)
        x = self.model[0](x)

        x, block_list = self.model[1](x)

        x_ori = x
        for i in range(5):
            x = block_list[i](x)
        x_fea1 = x

        for i in range(5):
            x = block_list[i + 5](x)
        x_fea2 = x

        for i in range(5):
            x = block_list[i + 10](x)
        x_fea3 = x

        for i in range(5):
            x = block_list[i + 15](x)
        x_fea4 = x

        x = block_list[20:](x)
        # short cut
        x = x_ori + x
        x = self.model[2:](x)
        x = self.HR_conv1_new(x)

        x_b_fea = self.b_fea_conv(x_grad)
        x_cat_1 = torch.cat([x_b_fea, x_fea1], dim=1)

        x_cat_1 = self.b_block_1(x_cat_1)
        x_cat_1 = self.b_concat_1(x_cat_1)

        x_cat_2 = torch.cat([x_cat_1, x_fea2], dim=1)

        x_cat_2 = self.b_block_2(x_cat_2)
        x_cat_2 = self.b_concat_2(x_cat_2)

        x_cat_3 = torch.cat([x_cat_2, x_fea3], dim=1)

        x_cat_3 = self.b_block_3(x_cat_3)
        x_cat_3 = self.b_concat_3(x_cat_3)

        x_cat_4 = torch.cat([x_cat_3, x_fea4], dim=1)

        x_cat_4 = self.b_block_4(x_cat_4)
        x_cat_4 = self.b_concat_4(x_cat_4)

        x_cat_4 = self.b_LR_conv(x_cat_4)

        # short cut
        x_cat_4 = x_cat_4 + x_b_fea
        x_branch = self.b_module(x_cat_4)

        # x_out_branch = self.conv_w(x_branch)
        ########
        x_branch_d = x_branch
        x_f_cat = torch.cat([x_branch_d, x], dim=1)
        x_f_cat = self.f_block(x_f_cat)
        x_out = self.f_concat(x_f_cat)
        x_out = self.f_HR_conv0(x_out)
        x_out = self.f_HR_conv1(x_out)

        #########
        # return x_out_branch, x_out, x_grad
        return x_out

nodes/impl/pytorch/architecture/SPSR.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# model: SRFormer
# SRFormer: Permuted Self-Attention for Single Image Super-Resolution

import math
import re

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint

from .timm.helpers import to_2tuple
from .timm.weight_init import trunc_normal_


class emptyModule(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x


def drop_path(x, drop_prob: float = 0.0, training: bool = False):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

    From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py
    """
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (
        x.ndim - 1
    )  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).

    From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py
    """

    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)


class dwconv(nn.Module):
    def __init__(self, hidden_features):
        super(dwconv, self).__init__()
        self.depthwise_conv = nn.Sequential(
            nn.Conv2d(
                hidden_features,
                hidden_features,
                kernel_size=5,
                stride=1,
                padding=2,
                dilation=1,
                groups=hidden_features,
            ),
            nn.GELU(),
        )
        self.hidden_features = hidden_features

    def forward(self, x, x_size):
        x = (
            x.transpose(1, 2)
            .view(x.shape[0], self.hidden_features, x_size[0], x_size[1])
            .contiguous()
        )  # b Ph*Pw c
        x = self.depthwise_conv(x)
        x = x.flatten(2).transpose(1, 2).contiguous()
        return x


class ConvFFN(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        drop=0.0,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.before_add = emptyModule()
        self.after_add = emptyModule()
        self.dwconv = dwconv(hidden_features=hidden_features)
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x, x_size):
        x = self.fc1(x)
        x = self.act(x)
        x = self.before_add(x)
        x = x + self.dwconv(x, x_size)
        x = self.after_add(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


def window_partition(x, window_size):
    """
    Args:
        x: (b, h, w, c)
        window_size (int): window size

    Returns:
        windows: (num_windows*b, window_size, window_size, c)
    """
    b, h, w, c = x.shape
    x = x.view(b, h // window_size, window_size, w // window_size, window_size, c)
    windows = (
        x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, c)
    )
    return windows


def window_reverse(windows, window_size, h, w):
    """
    Args:
        windows: (num_windows*b, window_size, window_size, c)
        window_size (int): Window size
        h (int): Height of image
        w (int): Width of image

    Returns:
        x: (b, h, w, c)
    """
    b = int(windows.shape[0] / (h * w / window_size / window_size))
    x = windows.view(
        b, h // window_size, w // window_size, window_size, window_size, -1
    )
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(b, h, w, -1)
    return x


class PSA(nn.Module):
    r"""Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.

    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
    """

    def __init__(
        self,
        dim,
        window_size,
        num_heads,
        qkv_bias=True,
        qk_scale=None,
        attn_drop=0.0,
        proj_drop=0.0,
    ):
        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wh, Ww
        self.permuted_window_size = (window_size[0] // 2, window_size[1] // 2)
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        # define a parameter table of relative position bias
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros(
                (2 * self.permuted_window_size[0] - 1)
                * (2 * self.permuted_window_size[1] - 1),
                num_heads,
            )
        )  # 2*Wh-1 * 2*Ww-1, nH

        # get pair-wise aligned relative position index for each token inside the window
        coords_h = torch.arange(self.permuted_window_size[0])
        coords_w = torch.arange(self.permuted_window_size[1])
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
        relative_coords = (
            coords_flatten[:, :, None] - coords_flatten[:, None, :]
        )  # 2, Wh*Ww, Wh*Ww
        relative_coords = relative_coords.permute(
            1, 2, 0
        ).contiguous()  # Wh*Ww, Wh*Ww, 2
        relative_coords[:, :, 0] += (
            self.permuted_window_size[0] - 1
        )  # shift to start from 0
        relative_coords[:, :, 1] += self.permuted_window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * self.permuted_window_size[1] - 1
        aligned_relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
        aligned_relative_position_index = (
            aligned_relative_position_index.reshape(
                self.permuted_window_size[0],
                self.permuted_window_size[1],
                1,
                1,
                self.permuted_window_size[0] * self.permuted_window_size[1],
            )
            .repeat(1, 1, 2, 2, 1)
            .permute(0, 2, 1, 3, 4)
            .reshape(
                4 * self.permuted_window_size[0] * self.permuted_window_size[1],
                self.permuted_window_size[0] * self.permuted_window_size[1],
            )
        )  #  FN*FN,WN*WN
        self.register_buffer(
            "aligned_relative_position_index", aligned_relative_position_index
        )
        # compresses the channel dimension of KV
        self.kv = nn.Linear(dim, dim // 2, bias=qkv_bias)
        self.q = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        trunc_normal_(self.relative_position_bias_table, std=0.02)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x, mask=None):
        """
        Args:
            x: input features with shape of (num_windows*b, n, c)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        b_, n, c = x.shape
        # compress the channel dimension of KV :(num_windows*b, num_heads, n//4, c//num_heads)
        kv = (
            self.kv(x)
            .reshape(
                b_,
                self.permuted_window_size[0],
                2,
                self.permuted_window_size[1],
                2,
                2,
                c // 4,
            )
            .permute(0, 1, 3, 5, 2, 4, 6)
            .reshape(b_, n // 4, 2, self.num_heads, c // self.num_heads)
            .permute(2, 0, 3, 1, 4)
        )
        k, v = kv[0], kv[1]
        # keep the channel dimension of Q: (num_windows*b, num_heads, n, c//num_heads)
        q = (
            self.q(x)
            .reshape(b_, n, 1, self.num_heads, c // self.num_heads)
            .permute(2, 0, 3, 1, 4)[0]
        )
        q = q * self.scale
        attn = q @ k.transpose(-2, -1)  # (num_windows*b, num_heads, n, n//4)

        relative_position_bias = self.relative_position_bias_table[
            self.aligned_relative_position_index.view(-1)
        ].view(
            self.window_size[0] * self.window_size[1],
            self.permuted_window_size[0] * self.permuted_window_size[1],
            -1,
        )  # (n, n//4)
        relative_position_bias = relative_position_bias.permute(
            2, 0, 1
        ).contiguous()  # (num_heads, n, n//4)
        attn = attn + relative_position_bias.unsqueeze(0)

        if mask is not None:
            nw = mask.shape[0]
            attn = attn.view(b_ // nw, nw, self.num_heads, n, n // 4) + mask.unsqueeze(
                1
            ).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, n, n // 4)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(b_, n, c)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

    def extra_repr(self) -> str:
        return f"dim={self.dim}, window_size={self.permuted_window_size}, num_heads={self.num_heads}"

    def flops(self, n):
        # calculate flops for 1 window with token length of n
        flops = 0
        # qkv = self.qkv(x)
        flops += n * self.dim * 1.5 * self.dim
        # attn = (q @ k.transpose(-2, -1))
        flops += self.num_heads * n * (self.dim // self.num_heads) * n / 4
        #  x = (attn @ v)
        flops += self.num_heads * n * n / 4 * (self.dim // self.num_heads)
        # x = self.proj(x)
        flops += n * self.dim * self.dim
        return flops


class PSA_Block(nn.Module):
    r"""Swin Transformer Block.

    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        num_heads (int): Number of attention heads.
        window_size (int): Window size.
        shift_size (int): Shift size for SW-MSA.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float, optional): Stochastic depth rate. Default: 0.0
        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(
        self,
        dim,
        input_resolution,
        num_heads,
        window_size=7,
        shift_size=0,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.permuted_window_size = window_size // 2
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        if min(self.input_resolution) <= self.window_size:
            # if window size is larger than input resolution, we don't partition windows
            self.shift_size = 0
            self.window_size = min(self.input_resolution)
        assert (
            0 <= self.shift_size < self.window_size
        ), "shift_size must in 0-window_size"
        self.norm1 = norm_layer(dim)

        self.attn = PSA(
            dim,
            window_size=to_2tuple(self.window_size),
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            attn_drop=attn_drop,
            proj_drop=drop,
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = ConvFFN(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

        if self.shift_size > 0:
            attn_mask = self.calculate_mask(self.input_resolution)
        else:
            attn_mask = None
        self.register_buffer("attn_mask", attn_mask)

        # emptyModule for Power Spectrum Based Evaluation
        self.after_norm1 = emptyModule()
        self.after_attention = emptyModule()
        self.residual_after_attention = emptyModule()
        self.after_norm2 = emptyModule()
        self.after_mlp = emptyModule()
        self.residual_after_mlp = emptyModule()

    def calculate_mask(self, x_size):
        # calculate mask for original windows
        h, w = x_size
        img_mask = torch.zeros((1, h, w, 1))  # 1 h w 1
        h_slices = (
            slice(0, -self.window_size),
            slice(-self.window_size, -self.shift_size),
            slice(-self.shift_size, None),
        )
        w_slices = (
            slice(0, -self.window_size),
            slice(-self.window_size, -self.shift_size),
            slice(-self.shift_size, None),
        )
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                img_mask[:, h, w, :] = cnt
                cnt += 1

        mask_windows = window_partition(
            img_mask, self.window_size
        )  # nw, window_size, window_size, 1
        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
        # calculate mask for permuted windows
        h, w = x_size
        permuted_window_mask = torch.zeros((1, h // 2, w // 2, 1))  # 1 h w 1
        h_slices = (
            slice(0, -self.permuted_window_size),
            slice(-self.permuted_window_size, -self.shift_size // 2),
            slice(-self.shift_size // 2, None),
        )
        w_slices = (
            slice(0, -self.permuted_window_size),
            slice(-self.permuted_window_size, -self.shift_size // 2),
            slice(-self.shift_size // 2, None),
        )
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                permuted_window_mask[:, h, w, :] = cnt
                cnt += 1

        permuted_windows = window_partition(
            permuted_window_mask, self.permuted_window_size
        )
        permuted_windows = permuted_windows.view(
            -1, self.permuted_window_size * self.permuted_window_size
        )
        # calculate attention mask
        attn_mask = mask_windows.unsqueeze(2) - permuted_windows.unsqueeze(1)
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
            attn_mask == 0, float(0.0)
        )

        return attn_mask

    def forward(self, x, x_size):
        h, w = x_size
        b, _, c = x.shape
        # assert seq_len == h * w, "input feature has wrong size"

        shortcut = x
        x = self.norm1(x)
        x = self.after_norm1(x)
        x = x.view(b, h, w, c)

        # cyclic shift
        if self.shift_size > 0:
            shifted_x = torch.roll(
                x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)
            )
        else:
            shifted_x = x

        # partition windows
        x_windows = window_partition(
            shifted_x, self.window_size
        )  # nw*b, window_size, window_size, c
        x_windows = x_windows.view(
            -1, self.window_size * self.window_size, c
        )  # nw*b, window_size*window_size, c

        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
        if self.input_resolution == x_size:
            attn_windows = self.attn(
                x_windows, mask=self.attn_mask
            )  # nw*b, window_size*window_size, c
        else:
            attn_windows = self.attn(
                x_windows, mask=self.calculate_mask(x_size).to(x.device)
            )

        # merge windows
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, c)
        shifted_x = window_reverse(attn_windows, self.window_size, h, w)  # b h' w' c

        # reverse cyclic shift
        if self.shift_size > 0:
            x = torch.roll(
                shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)
            )
        else:
            x = shifted_x
        x = x.view(b, h * w, c)
        x = self.after_attention(x)
        # FFN
        x = shortcut + self.drop_path(x)
        x = self.residual_after_attention(x)
        x = self.residual_after_mlp(
            x
            + self.drop_path(
                self.after_mlp(self.mlp(self.after_norm2(self.norm2(x)), x_size))
            )
        )

        return x

    def extra_repr(self) -> str:
        return (
            f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, "
            f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
        )

    def flops(self):
        flops = 0
        h, w = self.input_resolution
        # norm1
        flops += self.dim * h * w
        # W-MSA/SW-MSA
        nw = h * w / self.window_size / self.window_size
        flops += nw * self.attn.flops(self.window_size * self.window_size)
        # mlp
        flops += 2 * h * w * self.dim * self.dim * self.mlp_ratio
        flops += h * w * self.dim * 25
        # norm2
        flops += self.dim * h * w
        return flops


class PatchMerging(nn.Module):
    r"""Patch Merging Layer.

    Args:
        input_resolution (tuple[int]): Resolution of input feature.
        dim (int): Number of input channels.
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(4 * dim)

    def forward(self, x):
        """
        x: b, h*w, c
        """
        h, w = self.input_resolution
        b, seq_len, c = x.shape
        assert seq_len == h * w, "input feature has wrong size"
        assert h % 2 == 0 and w % 2 == 0, f"x size ({h}*{w}) are not even."

        x = x.view(b, h, w, c)

        x0 = x[:, 0::2, 0::2, :]  # b h/2 w/2 c
        x1 = x[:, 1::2, 0::2, :]  # b h/2 w/2 c
        x2 = x[:, 0::2, 1::2, :]  # b h/2 w/2 c
        x3 = x[:, 1::2, 1::2, :]  # b h/2 w/2 c
        x = torch.cat([x0, x1, x2, x3], -1)  # b h/2 w/2 4*c
        x = x.view(b, -1, 4 * c)  # b h/2*w/2 4*c

        x = self.norm(x)
        x = self.reduction(x)

        return x

    def extra_repr(self) -> str:
        return f"input_resolution={self.input_resolution}, dim={self.dim}"

    def flops(self):
        h, w = self.input_resolution
        flops = h * w * self.dim
        flops += (h // 2) * (w // 2) * 4 * self.dim * 2 * self.dim
        return flops


class BasicLayer(nn.Module):
    """A basic Swin Transformer layer for one stage.

    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
    """

    def __init__(
        self,
        dim,
        input_resolution,
        depth,
        num_heads,
        window_size,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        downsample=None,
        use_checkpoint=False,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth
        self.use_checkpoint = use_checkpoint

        # build blocks
        self.blocks = nn.ModuleList(
            [
                PSA_Block(
                    dim=dim,
                    input_resolution=input_resolution,
                    num_heads=num_heads,
                    window_size=window_size,
                    shift_size=0 if (i % 2 == 0) else window_size // 2,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop,
                    attn_drop=attn_drop,
                    drop_path=drop_path[i]
                    if isinstance(drop_path, list)
                    else drop_path,
                    norm_layer=norm_layer,
                )
                for i in range(depth)
            ]
        )

        # patch merging layer
        if downsample is not None:
            self.downsample = downsample(
                input_resolution, dim=dim, norm_layer=norm_layer
            )
        else:
            self.downsample = None

    def forward(self, x, x_size):
        for blk in self.blocks:
            if self.use_checkpoint:
                x = checkpoint.checkpoint(blk, x)
            else:
                x = blk(x, x_size)
        if self.downsample is not None:
            x = self.downsample(x)
        return x

    def extra_repr(self) -> str:
        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"

    def flops(self):
        flops = 0
        for blk in self.blocks:
            flops += blk.flops()
        if self.downsample is not None:
            flops += self.downsample.flops()
        return flops


class PSA_Group(nn.Module):
    """Residual Swin Transformer Block (PSA_Group).

    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
        img_size: Input image size.
        patch_size: Patch size.
        resi_connection: The convolutional block before residual connection.
    """

    def __init__(
        self,
        dim,
        input_resolution,
        depth,
        num_heads,
        window_size,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        downsample=None,
        use_checkpoint=False,
        img_size=224,
        patch_size=4,
        resi_connection="1conv",
    ):
        super(PSA_Group, self).__init__()

        self.dim = dim
        self.input_resolution = input_resolution

        self.residual_group = BasicLayer(
            dim=dim,
            input_resolution=input_resolution,
            depth=depth,
            num_heads=num_heads,
            window_size=window_size,
            mlp_ratio=mlp_ratio,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            drop=drop,
            attn_drop=attn_drop,
            drop_path=drop_path,
            norm_layer=norm_layer,
            downsample=downsample,
            use_checkpoint=use_checkpoint,
        )

        if resi_connection == "1conv":
            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
        elif resi_connection == "3conv":
            # to save parameters and memory
            self.conv = nn.Sequential(
                nn.Conv2d(dim, dim // 4, 3, 1, 1),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(dim // 4, dim, 3, 1, 1),
            )

        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=0,
            embed_dim=dim,
            norm_layer=None,
        )

        self.patch_unembed = PatchUnEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=0,
            embed_dim=dim,
            norm_layer=None,
        )

        self.before_PSA_Group_conv = emptyModule()
        self.after_PSA_Group_conv = emptyModule()
        self.after_PSA_Group_Residual = emptyModule()

    def forward(self, x, x_size):
        return self.after_PSA_Group_Residual(
            self.after_PSA_Group_conv(
                self.patch_embed(
                    self.conv(
                        self.patch_unembed(
                            self.before_PSA_Group_conv(self.residual_group(x, x_size)),
                            x_size,
                        )
                    )
                )
            )
            + x
        )

    def flops(self):
        flops = 0
        flops += self.residual_group.flops()
        h, w = self.input_resolution
        flops += h * w * self.dim * self.dim * 9
        flops += self.patch_embed.flops()
        flops += self.patch_unembed.flops()

        return flops


class PatchEmbed(nn.Module):
    r"""Image to Patch Embedding

    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
        self,
        img_size=224,
        window_size=22,
        patch_size=4,
        in_chans=3,
        embed_dim=96,
        norm_layer=None,
    ):
        super().__init__()
        if img_size % window_size != 0:
            img_size = img_size + (window_size - img_size % window_size)
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [
            img_size[0] // patch_size[0],
            img_size[1] // patch_size[1],
        ]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]

        self.in_chans = in_chans
        self.embed_dim = embed_dim

        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        x = x.flatten(2).transpose(1, 2)  # b Ph*Pw c
        if self.norm is not None:
            x = self.norm(x)
        return x

    def flops(self):
        flops = 0
        h, w = self.img_size
        if self.norm is not None:
            flops += h * w * self.embed_dim
        return flops


class PatchUnEmbed(nn.Module):
    r"""Image to Patch Unembedding

    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
        self,
        img_size=224,
        window_size=24,
        patch_size=4,
        in_chans=3,
        embed_dim=96,
        norm_layer=None,
    ):
        super().__init__()
        if img_size % window_size != 0:
            img_size = img_size + (window_size - img_size % window_size)
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [
            img_size[0] // patch_size[0],
            img_size[1] // patch_size[1],
        ]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]
        self.in_chans = in_chans
        self.embed_dim = embed_dim

    def forward(self, x, x_size):
        x = x.transpose(1, 2).view(
            x.shape[0], self.embed_dim, x_size[0], x_size[1]
        )  # b Ph*Pw c
        return x

    def flops(self):
        flops = 0
        return flops


class Upsample(nn.Sequential):
    """Upsample module.

    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.
    """

    def __init__(self, scale, num_feat):
        m = []
        if (scale & (scale - 1)) == 0:  # scale = 2^n
            for _ in range(int(math.log(scale, 2))):
                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
                m.append(nn.PixelShuffle(2))
        elif scale == 3:
            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
            m.append(nn.PixelShuffle(3))
        else:
            raise ValueError(
                f"scale {scale} is not supported. Supported scales: 2^n and 3."
            )
        super(Upsample, self).__init__(*m)


class UpsampleOneStep(nn.Sequential):
    """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
       Used in lightweight SR to save parameters.

    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.

    """

    def __init__(self, scale, num_feat, num_out_ch, input_resolution=None):
        self.num_feat = num_feat
        self.input_resolution = input_resolution
        m = []
        m.append(nn.Conv2d(num_feat, (scale**2) * num_out_ch, 3, 1, 1))
        m.append(nn.PixelShuffle(scale))
        super(UpsampleOneStep, self).__init__(*m)

    def flops(self):
        h, w = self.input_resolution
        flops = h * w * self.num_feat * 3 * 9
        return flops


class SRFormer(nn.Module):
    r"""SRFormer
        A PyTorch implement of : `SRFormer: Permuted Self-Attention for Single Image Super-Resolution`, based on Swin Transformer.

    Args:
        img_size (int | tuple(int)): Input image size. Default 64
        patch_size (int | tuple(int)): Patch size. Default: 1
        in_chans (int): Number of input image channels. Default: 3
        embed_dim (int): Patch embedding dimension. Default: 96
        depths (tuple(int)): Depth of each Swin Transformer layer.
        num_heads (tuple(int)): Number of attention heads in different layers.
        window_size (int): Window size. Default: 7
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
        drop_rate (float): Dropout rate. Default: 0
        attn_drop_rate (float): Attention dropout rate. Default: 0
        drop_path_rate (float): Stochastic depth rate. Default: 0.1
        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
        patch_norm (bool): If True, add normalization after patch embedding. Default: True
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
        upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
        img_range: Image range. 1. or 255.
        upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
    """

    def __init__(self, state_dict):
        super(SRFormer, self).__init__()

        # Default
        img_size = 64
        patch_size = 1
        in_chans = 3
        embed_dim = 96
        depths = (6, 6, 6, 6)
        num_heads = (6, 6, 6, 6)
        window_size = 7
        mlp_ratio = 4.0
        qkv_bias = True
        qk_scale = None
        drop_rate = 0.0
        attn_drop_rate = 0.0
        drop_path_rate = 0.1
        norm_layer = nn.LayerNorm
        ape = False
        patch_norm = True
        use_checkpoint = False
        upscale = 2
        img_range = 1.0
        upsampler = ""
        resi_connection = "1conv"

        self.model_arch = "SRFormer"
        self.sub_type = "SR"
        self.state = state_dict

        state_keys = list(state_dict.keys())

        if "conv_before_upsample.0.weight" in state_keys:
            if "conv_up1.weight" in state_keys:
                upsampler = "nearest+conv"
            else:
                upsampler = "pixelshuffle"
        elif "upsample.0.weight" in state_keys:
            upsampler = "pixelshuffledirect"
        else:
            upsampler = ""

        num_feat = (
            state_dict.get("conv_before_upsample.0.weight", None).shape[1]
            if state_dict.get("conv_before_upsample.weight", None)
            else 64
        )

        num_in_ch = state_dict["conv_first.weight"].shape[1]
        in_chans = num_in_ch
        if "conv_last.weight" in state_keys:
            num_out_ch = state_dict["conv_last.weight"].shape[0]
        else:
            num_out_ch = num_in_ch

        upscale = 1
        if upsampler == "nearest+conv":
            upsample_keys = [
                x for x in state_keys if "conv_up" in x and "bias" not in x
            ]

            for upsample_key in upsample_keys:
                upscale *= 2
        elif upsampler == "pixelshuffle":
            upsample_keys = [
                x
                for x in state_keys
                if "upsample" in x and "conv" not in x and "bias" not in x
            ]
            for upsample_key in upsample_keys:
                shape = state_dict[upsample_key].shape[0]
                upscale *= math.sqrt(shape // num_feat)
            upscale = int(upscale)
        elif upsampler == "pixelshuffledirect":
            upscale = int(
                math.sqrt(state_dict["upsample.0.bias"].shape[0] // num_out_ch)
            )

        max_layer_num = 0
        max_block_num = 0
        for key in state_keys:
            result = re.match(
                r"layers.(\d*).residual_group.blocks.(\d*).norm1.weight", key
            )
            if result:
                layer_num, block_num = result.groups()
                max_layer_num = max(max_layer_num, int(layer_num))
                max_block_num = max(max_block_num, int(block_num))

        depths = [max_block_num + 1 for _ in range(max_layer_num + 1)]

        if (
            "layers.0.residual_group.blocks.0.attn.relative_position_bias_table"
            in state_keys
        ):
            num_heads_num = state_dict[
                "layers.0.residual_group.blocks.0.attn.relative_position_bias_table"
            ].shape[-1]
            num_heads = [num_heads_num for _ in range(max_layer_num + 1)]
        else:
            num_heads = depths

        embed_dim = state_dict["conv_first.weight"].shape[0]

        mlp_ratio = float(
            state_dict["layers.0.residual_group.blocks.0.mlp.fc1.bias"].shape[0]
            / embed_dim
        )

        # TODO: could actually count the layers, but this should do
        # TOOD: confirm this is correct and the same as SwinIR
        if "layers.0.conv.4.weight" in state_keys:
            resi_connection = "3conv"
        else:
            resi_connection = "1conv"

        window_size = int(
            math.sqrt(
                state_dict[
                    "layers.0.residual_group.blocks.0.attn.aligned_relative_position_index"
                ].shape[0]
            )
        )

        if "layers.0.residual_group.blocks.1.attn_mask" in state_keys:
            img_size = int(
                (
                    math.sqrt(
                        state_dict["layers.0.residual_group.blocks.1.attn_mask"].shape[
                            0
                        ]
                    )
                )
                * 16
            )

        self.in_nc = num_in_ch
        self.out_nc = num_out_ch
        self.num_feat = num_feat
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.depths = depths
        self.window_size = window_size
        self.mlp_ratio = mlp_ratio
        self.scale = upscale
        self.upsampler = upsampler
        self.img_size = img_size
        self.img_range = img_range
        self.resi_connection = resi_connection

        self.supports_fp16 = False  # Too much weirdness to support this at the moment
        self.supports_bfp16 = True
        self.min_size_restriction = 16

        num_in_ch = in_chans
        num_out_ch = in_chans
        num_feat = 64
        self.img_range = img_range
        if in_chans == 3:
            rgb_mean = (0.4488, 0.4371, 0.4040)
            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
        else:
            self.mean = torch.zeros(1, 1, 1, 1)
        self.upscale = upscale
        self.upsampler = upsampler
        self.window_size = window_size
        # ------------------------- 1, shallow feature extraction ------------------------- #
        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)

        # ------------------------- 2, deep feature extraction ------------------------- #
        self.num_layers = len(depths)
        self.embed_dim = embed_dim
        self.ape = ape
        self.patch_norm = patch_norm
        self.num_features = embed_dim
        self.mlp_ratio = mlp_ratio

        # split image into non-overlapping patches
        self.patch_embed = PatchEmbed(
            img_size=img_size,
            window_size=window_size,
            patch_size=patch_size,
            in_chans=embed_dim,
            embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None,
        )
        num_patches = self.patch_embed.num_patches
        patches_resolution = self.patch_embed.patches_resolution
        self.patches_resolution = patches_resolution

        # merge non-overlapping patches into image
        self.patch_unembed = PatchUnEmbed(
            img_size=img_size,
            window_size=window_size,
            patch_size=patch_size,
            in_chans=embed_dim,
            embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None,
        )

        # absolute position embedding
        if self.ape:
            self.absolute_pos_embed = nn.Parameter(
                torch.zeros(1, num_patches, embed_dim)
            )
            trunc_normal_(self.absolute_pos_embed, std=0.02)

        self.pos_drop = nn.Dropout(p=drop_rate)

        # stochastic depth
        dpr = [
            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
        ]  # stochastic depth decay rule

        # build Permuted Self Attention Group  (PSA_Group)
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):
            layer = PSA_Group(
                dim=embed_dim,
                input_resolution=(patches_resolution[0], patches_resolution[1]),
                depth=depths[i_layer],
                num_heads=num_heads[i_layer],
                window_size=window_size,
                mlp_ratio=self.mlp_ratio,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop=drop_rate,
                attn_drop=attn_drop_rate,
                drop_path=dpr[
                    sum(depths[:i_layer]) : sum(depths[: i_layer + 1])
                ],  # no impact on SR results
                norm_layer=norm_layer,
                downsample=None,
                use_checkpoint=use_checkpoint,
                img_size=img_size,
                patch_size=patch_size,
                resi_connection=resi_connection,
            )
            self.layers.append(layer)
        self.norm = norm_layer(self.num_features)

        # build the last conv layer in deep feature extraction
        if resi_connection == "1conv":
            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
        elif resi_connection == "3conv":
            # to save parameters and memory
            self.conv_after_body = nn.Sequential(
                nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1),
            )

        # ------------------------- 3, high quality image reconstruction ------------------------- #
        if self.upsampler == "pixelshuffle":
            # for classical SR
            self.conv_before_upsample = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.upsample = Upsample(upscale, num_feat)
            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
        elif self.upsampler == "pixelshuffledirect":
            # for lightweight SR (to save parameters)
            self.upsample = UpsampleOneStep(
                upscale,
                embed_dim,
                num_out_ch,
                (patches_resolution[0], patches_resolution[1]),
            )
        elif self.upsampler == "nearest+conv":
            # for real-world SR (less artifacts)
            assert self.upscale == 4, "only support x4 now."
            self.conv_before_upsample = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
            self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
            self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
            self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        else:
            # for image denoising and JPEG compression artifact reduction
            self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1)

        self.apply(self._init_weights)

        self.load_state_dict(state_dict, strict=True)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {"absolute_pos_embed"}

    @torch.jit.ignore
    def no_weight_decay_keywords(self):
        return {"relative_position_bias_table"}

    def forward_features(self, x):
        x_size = (x.shape[2], x.shape[3])
        x = self.patch_embed(x)
        if self.ape:
            x = x + self.absolute_pos_embed
        x = self.pos_drop(x)

        for layer in self.layers:
            x = layer(x, x_size)

        x = self.norm(x)  # b seq_len c
        x = self.patch_unembed(x, x_size)

        return x

    def check_image_size(self, x):
        _, _, h, w = x.size()
        mod_pad_h = (self.window_size - h % self.window_size) % self.window_size
        mod_pad_w = (self.window_size - w % self.window_size) % self.window_size
        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect")
        return x

    def forward(self, x):
        H, W = x.shape[2:]
        x = self.check_image_size(x)
        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == "pixelshuffle":
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == "pixelshuffledirect":
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == "nearest+conv":
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(
                self.conv_up1(
                    torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest")
                )
            )
            x = self.lrelu(
                self.conv_up2(
                    torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest")
                )
            )
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean

        return x[:, :, : H * self.upscale, : W * self.upscale]

    def flops(self):
        flops = 0
        h, w = self.patches_resolution
        flops += h * w * 3 * self.embed_dim * 9
        flops += self.patch_embed.flops()
        for layer in self.layers:
            flops += layer.flops()
        flops += h * w * 9 * self.embed_dim * self.embed_dim
        flops += self.upsample.flops()
        return flops

nodes/impl/pytorch/architecture/SRFormer.py
vladmandic/sd-extension-chainner
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import math

import torch.nn as nn
import torch.nn.functional as F


class SRVGGNetCompact(nn.Module):
    """A compact VGG-style network structure for super-resolution.
    It is a compact network structure, which performs upsampling in the last layer and no convolution is
    conducted on the HR feature space.
    Args:
        num_in_ch (int): Channel number of inputs. Default: 3.
        num_out_ch (int): Channel number of outputs. Default: 3.
        num_feat (int): Channel number of intermediate features. Default: 64.
        num_conv (int): Number of convolution layers in the body network. Default: 16.
        upscale (int): Upsampling factor. Default: 4.
        act_type (str): Activation type, options: 'relu', 'prelu', 'leakyrelu'. Default: prelu.
    """

    def __init__(
        self,
        state_dict,
        act_type: str = "prelu",
    ):
        super(SRVGGNetCompact, self).__init__()
        self.model_arch = "SRVGG (RealESRGAN)"
        self.sub_type = "SR"

        self.act_type = act_type

        self.state = state_dict

        if "params" in self.state:
            self.state = self.state["params"]

        self.key_arr = list(self.state.keys())

        self.in_nc = self.get_in_nc()
        self.num_feat = self.get_num_feats()
        self.num_conv = self.get_num_conv()
        self.out_nc = self.in_nc  # :(
        self.pixelshuffle_shape = None  # Defined in get_scale()
        self.scale = self.get_scale()

        self.supports_fp16 = True
        self.supports_bfp16 = True
        self.min_size_restriction = None

        self.body = nn.ModuleList()
        # the first conv
        self.body.append(nn.Conv2d(self.in_nc, self.num_feat, 3, 1, 1))
        # the first activation
        if act_type == "relu":
            activation = nn.ReLU(inplace=True)
        elif act_type == "prelu":
            activation = nn.PReLU(num_parameters=self.num_feat)
        elif act_type == "leakyrelu":
            activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        self.body.append(activation)  # type: ignore

        # the body structure
        for _ in range(self.num_conv):
            self.body.append(nn.Conv2d(self.num_feat, self.num_feat, 3, 1, 1))
            # activation
            if act_type == "relu":
                activation = nn.ReLU(inplace=True)
            elif act_type == "prelu":
                activation = nn.PReLU(num_parameters=self.num_feat)
            elif act_type == "leakyrelu":
                activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
            self.body.append(activation)  # type: ignore

        # the last conv
        self.body.append(nn.Conv2d(self.num_feat, self.pixelshuffle_shape, 3, 1, 1))  # type: ignore
        # upsample
        self.upsampler = nn.PixelShuffle(self.scale)

        self.load_state_dict(self.state, strict=False)

    def get_num_conv(self) -> int:
        return (int(self.key_arr[-1].split(".")[1]) - 2) // 2

    def get_num_feats(self) -> int:
        return self.state[self.key_arr[0]].shape[0]

    def get_in_nc(self) -> int:
        return self.state[self.key_arr[0]].shape[1]

    def get_scale(self) -> int:
        self.pixelshuffle_shape = self.state[self.key_arr[-1]].shape[0]
        # Assume out_nc is the same as in_nc
        # I cant think of a better way to do that
        self.out_nc = self.in_nc
        scale = math.sqrt(self.pixelshuffle_shape / self.out_nc)
        if scale - int(scale) > 0:
            print(
                "out_nc is probably different than in_nc, scale calculation might be wrong"
            )
        scale = int(scale)
        return scale

    def forward(self, x):
        out = x
        for i in range(0, len(self.body)):
            out = self.body[i](out)

        out = self.upsampler(out)
        # add the nearest upsampled image, so that the network learns the residual
        base = F.interpolate(x, scale_factor=self.scale, mode="nearest")
        out += base
        return out

nodes/impl/pytorch/architecture/SRVGG.py
vladmandic/sd-extension-chainner
# From https://github.com/Koushik0901/Swift-SRGAN/blob/master/swift-srgan/models.py

import torch
from torch import nn


class SeperableConv2d(nn.Module):
    def __init__(
        self, in_channels, out_channels, kernel_size, stride=1, padding=1, bias=True
    ):
        super(SeperableConv2d, self).__init__()
        self.depthwise = nn.Conv2d(
            in_channels,
            in_channels,
            kernel_size=kernel_size,
            stride=stride,
            groups=in_channels,
            bias=bias,
            padding=padding,
        )
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)

    def forward(self, x):
        return self.pointwise(self.depthwise(x))


class ConvBlock(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        use_act=True,
        use_bn=True,
        discriminator=False,
        **kwargs,
    ):
        super(ConvBlock, self).__init__()

        self.use_act = use_act
        self.cnn = SeperableConv2d(in_channels, out_channels, **kwargs, bias=not use_bn)
        self.bn = nn.BatchNorm2d(out_channels) if use_bn else nn.Identity()
        self.act = (
            nn.LeakyReLU(0.2, inplace=True)
            if discriminator
            else nn.PReLU(num_parameters=out_channels)
        )

    def forward(self, x):
        return self.act(self.bn(self.cnn(x))) if self.use_act else self.bn(self.cnn(x))


class UpsampleBlock(nn.Module):
    def __init__(self, in_channels, scale_factor):
        super(UpsampleBlock, self).__init__()

        self.conv = SeperableConv2d(
            in_channels,
            in_channels * scale_factor**2,
            kernel_size=3,
            stride=1,
            padding=1,
        )
        self.ps = nn.PixelShuffle(
            scale_factor
        )  # (in_channels * 4, H, W) -> (in_channels, H*2, W*2)
        self.act = nn.PReLU(num_parameters=in_channels)

    def forward(self, x):
        return self.act(self.ps(self.conv(x)))


class ResidualBlock(nn.Module):
    def __init__(self, in_channels):
        super(ResidualBlock, self).__init__()

        self.block1 = ConvBlock(
            in_channels, in_channels, kernel_size=3, stride=1, padding=1
        )
        self.block2 = ConvBlock(
            in_channels, in_channels, kernel_size=3, stride=1, padding=1, use_act=False
        )

    def forward(self, x):
        out = self.block1(x)
        out = self.block2(out)
        return out + x


class Generator(nn.Module):
    """Swift-SRGAN Generator
    Args:
        in_channels (int): number of input image channels.
        num_channels (int): number of hidden channels.
        num_blocks (int): number of residual blocks.
        upscale_factor (int): factor to upscale the image [2x, 4x, 8x].
    Returns:
        torch.Tensor: super resolution image
    """

    def __init__(
        self,
        state_dict,
    ):
        super(Generator, self).__init__()
        self.model_arch = "Swift-SRGAN"
        self.sub_type = "SR"
        self.state = state_dict
        if "model" in self.state:
            self.state = self.state["model"]

        self.in_nc: int = self.state["initial.cnn.depthwise.weight"].shape[0]
        self.out_nc: int = self.state["final_conv.pointwise.weight"].shape[0]
        self.num_filters: int = self.state["initial.cnn.pointwise.weight"].shape[0]
        self.num_blocks = len(
            {x.split(".")[1] for x in self.state.keys() if "residual" in x}
        )
        self.scale: int = 2 ** len(
            {x.split(".")[1] for x in self.state.keys() if "upsampler" in x}
        )

        in_channels = self.in_nc
        num_channels = self.num_filters
        num_blocks = self.num_blocks
        upscale_factor = self.scale

        self.supports_fp16 = True
        self.supports_bfp16 = True
        self.min_size_restriction = None

        self.initial = ConvBlock(
            in_channels, num_channels, kernel_size=9, stride=1, padding=4, use_bn=False
        )
        self.residual = nn.Sequential(
            *[ResidualBlock(num_channels) for _ in range(num_blocks)]
        )
        self.convblock = ConvBlock(
            num_channels,
            num_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            use_act=False,
        )
        self.upsampler = nn.Sequential(
            *[
                UpsampleBlock(num_channels, scale_factor=2)
                for _ in range(upscale_factor // 2)
            ]
        )
        self.final_conv = SeperableConv2d(
            num_channels, in_channels, kernel_size=9, stride=1, padding=4
        )

        self.load_state_dict(self.state, strict=False)

    def forward(self, x):
        initial = self.initial(x)
        x = self.residual(initial)
        x = self.convblock(x) + initial
        x = self.upsampler(x)
        return (torch.tanh(self.final_conv(x)) + 1) / 2

nodes/impl/pytorch/architecture/SwiftSRGAN.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# -----------------------------------------------------------------------------------
# Swin2SR: Swin2SR: SwinV2 Transformer for Compressed Image Super-Resolution and Restoration, https://arxiv.org/abs/2209.11345
# Written by Conde and Choi et al.
# From: https://raw.githubusercontent.com/mv-lab/swin2sr/main/models/network_swin2sr.py
# -----------------------------------------------------------------------------------

import math
import re

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint

# Originally from the timm package
from .timm.drop import DropPath
from .timm.helpers import to_2tuple
from .timm.weight_init import trunc_normal_


class Mlp(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        drop=0.0,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


def window_partition(x, window_size):
    """
    Args:
        x: (B, H, W, C)
        window_size (int): window size
    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = (
        x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    )
    return windows


def window_reverse(windows, window_size, H, W):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size (int): Window size
        H (int): Height of image
        W (int): Width of image
    Returns:
        x: (B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(
        B, H // window_size, W // window_size, window_size, window_size, -1
    )
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x


class WindowAttention(nn.Module):
    r"""Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.
    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
        pretrained_window_size (tuple[int]): The height and width of the window in pre-training.
    """

    def __init__(
        self,
        dim,
        window_size,
        num_heads,
        qkv_bias=True,
        attn_drop=0.0,
        proj_drop=0.0,
        pretrained_window_size=None,
    ):
        if pretrained_window_size is None:
            pretrained_window_size = [0, 0]
        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wh, Ww
        self.pretrained_window_size = pretrained_window_size
        self.num_heads = num_heads

        self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)  # type: ignore

        # mlp to generate continuous relative position bias
        self.cpb_mlp = nn.Sequential(
            nn.Linear(2, 512, bias=True),
            nn.ReLU(inplace=True),
            nn.Linear(512, num_heads, bias=False),
        )

        # get relative_coords_table
        relative_coords_h = torch.arange(
            -(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32
        )
        relative_coords_w = torch.arange(
            -(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32
        )
        relative_coords_table = (
            torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w]))
            .permute(1, 2, 0)
            .contiguous()
            .unsqueeze(0)
        )  # 1, 2*Wh-1, 2*Ww-1, 2
        if pretrained_window_size[0] > 0:
            relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1
            relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1
        else:
            relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1
            relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1
        relative_coords_table *= 8  # normalize to -8, 8
        relative_coords_table = (
            torch.sign(relative_coords_table)
            * torch.log2(torch.abs(relative_coords_table) + 1.0)
            / np.log2(8)
        )

        self.register_buffer("relative_coords_table", relative_coords_table)

        # get pair-wise relative position index for each token inside the window
        coords_h = torch.arange(self.window_size[0])
        coords_w = torch.arange(self.window_size[1])
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
        relative_coords = (
            coords_flatten[:, :, None] - coords_flatten[:, None, :]
        )  # 2, Wh*Ww, Wh*Ww
        relative_coords = relative_coords.permute(
            1, 2, 0
        ).contiguous()  # Wh*Ww, Wh*Ww, 2
        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
        self.register_buffer("relative_position_index", relative_position_index)

        self.qkv = nn.Linear(dim, dim * 3, bias=False)
        if qkv_bias:
            self.q_bias = nn.Parameter(torch.zeros(dim))  # type: ignore
            self.v_bias = nn.Parameter(torch.zeros(dim))  # type: ignore
        else:
            self.q_bias = None
            self.v_bias = None
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x, mask=None):
        """
        Args:
            x: input features with shape of (num_windows*B, N, C)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        B_, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))  # type: ignore
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = (
            qkv[0],
            qkv[1],
            qkv[2],
        )  # make torchscript happy (cannot use tensor as tuple)

        # cosine attention
        attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)
        logit_scale = torch.clamp(
            self.logit_scale,
            max=torch.log(torch.tensor(1.0 / 0.01)).to(self.logit_scale.device),
        ).exp()
        attn = attn * logit_scale

        relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(
            -1, self.num_heads
        )
        relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(  # type: ignore
            self.window_size[0] * self.window_size[1],
            self.window_size[0] * self.window_size[1],
            -1,
        )  # Wh*Ww,Wh*Ww,nH
        relative_position_bias = relative_position_bias.permute(
            2, 0, 1
        ).contiguous()  # nH, Wh*Ww, Wh*Ww
        relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
        attn = attn + relative_position_bias.unsqueeze(0)

        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(
                1
            ).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, N, N)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

    def extra_repr(self) -> str:
        return (
            f"dim={self.dim}, window_size={self.window_size}, "
            f"pretrained_window_size={self.pretrained_window_size}, num_heads={self.num_heads}"
        )

    def flops(self, N):
        # calculate flops for 1 window with token length of N
        flops = 0
        # qkv = self.qkv(x)
        flops += N * self.dim * 3 * self.dim
        # attn = (q @ k.transpose(-2, -1))
        flops += self.num_heads * N * (self.dim // self.num_heads) * N
        #  x = (attn @ v)
        flops += self.num_heads * N * N * (self.dim // self.num_heads)
        # x = self.proj(x)
        flops += N * self.dim * self.dim
        return flops


class SwinTransformerBlock(nn.Module):
    r"""Swin Transformer Block.
    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resulotion.
        num_heads (int): Number of attention heads.
        window_size (int): Window size.
        shift_size (int): Shift size for SW-MSA.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float, optional): Stochastic depth rate. Default: 0.0
        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
        pretrained_window_size (int): Window size in pre-training.
    """

    def __init__(
        self,
        dim,
        input_resolution,
        num_heads,
        window_size=7,
        shift_size=0,
        mlp_ratio=4.0,
        qkv_bias=True,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        pretrained_window_size=0,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        if min(self.input_resolution) <= self.window_size:
            # if window size is larger than input resolution, we don't partition windows
            self.shift_size = 0
            self.window_size = min(self.input_resolution)
        assert (
            0 <= self.shift_size < self.window_size
        ), "shift_size must in 0-window_size"

        self.norm1 = norm_layer(dim)
        self.attn = WindowAttention(
            dim,
            window_size=to_2tuple(self.window_size),
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
            pretrained_window_size=to_2tuple(pretrained_window_size),
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

        if self.shift_size > 0:
            attn_mask = self.calculate_mask(self.input_resolution)
        else:
            attn_mask = None

        self.register_buffer("attn_mask", attn_mask)

    def calculate_mask(self, x_size):
        # calculate attention mask for SW-MSA
        H, W = x_size
        img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
        h_slices = (
            slice(0, -self.window_size),
            slice(-self.window_size, -self.shift_size),
            slice(-self.shift_size, None),
        )
        w_slices = (
            slice(0, -self.window_size),
            slice(-self.window_size, -self.shift_size),
            slice(-self.shift_size, None),
        )
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                img_mask[:, h, w, :] = cnt
                cnt += 1

        mask_windows = window_partition(
            img_mask, self.window_size
        )  # nW, window_size, window_size, 1
        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
            attn_mask == 0, float(0.0)
        )

        return attn_mask

    def forward(self, x, x_size):
        H, W = x_size
        B, L, C = x.shape
        # assert L == H * W, "input feature has wrong size"

        shortcut = x
        x = x.view(B, H, W, C)

        # cyclic shift
        if self.shift_size > 0:
            shifted_x = torch.roll(
                x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)
            )
        else:
            shifted_x = x

        # partition windows
        x_windows = window_partition(
            shifted_x, self.window_size
        )  # nW*B, window_size, window_size, C
        x_windows = x_windows.view(
            -1, self.window_size * self.window_size, C
        )  # nW*B, window_size*window_size, C

        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
        if self.input_resolution == x_size:
            attn_windows = self.attn(
                x_windows, mask=self.attn_mask
            )  # nW*B, window_size*window_size, C
        else:
            attn_windows = self.attn(
                x_windows, mask=self.calculate_mask(x_size).to(x.device)
            )

        # merge windows
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C

        # reverse cyclic shift
        if self.shift_size > 0:
            x = torch.roll(
                shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)
            )
        else:
            x = shifted_x
        x = x.view(B, H * W, C)
        x = shortcut + self.drop_path(self.norm1(x))

        # FFN
        x = x + self.drop_path(self.norm2(self.mlp(x)))

        return x

    def extra_repr(self) -> str:
        return (
            f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, "
            f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
        )

    def flops(self):
        flops = 0
        H, W = self.input_resolution
        # norm1
        flops += self.dim * H * W
        # W-MSA/SW-MSA
        nW = H * W / self.window_size / self.window_size
        flops += nW * self.attn.flops(self.window_size * self.window_size)
        # mlp
        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
        # norm2
        flops += self.dim * H * W
        return flops


class PatchMerging(nn.Module):
    r"""Patch Merging Layer.
    Args:
        input_resolution (tuple[int]): Resolution of input feature.
        dim (int): Number of input channels.
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(2 * dim)

    def forward(self, x):
        """
        x: B, H*W, C
        """
        H, W = self.input_resolution
        B, L, C = x.shape
        assert L == H * W, "input feature has wrong size"
        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."

        x = x.view(B, H, W, C)

        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C

        x = self.reduction(x)
        x = self.norm(x)

        return x

    def extra_repr(self) -> str:
        return f"input_resolution={self.input_resolution}, dim={self.dim}"

    def flops(self):
        H, W = self.input_resolution
        flops = (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
        flops += H * W * self.dim // 2
        return flops


class BasicLayer(nn.Module):
    """A basic Swin Transformer layer for one stage.
    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
        pretrained_window_size (int): Local window size in pre-training.
    """

    def __init__(
        self,
        dim,
        input_resolution,
        depth,
        num_heads,
        window_size,
        mlp_ratio=4.0,
        qkv_bias=True,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        downsample=None,
        use_checkpoint=False,
        pretrained_window_size=0,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth
        self.use_checkpoint = use_checkpoint

        # build blocks
        self.blocks = nn.ModuleList(
            [
                SwinTransformerBlock(
                    dim=dim,
                    input_resolution=input_resolution,
                    num_heads=num_heads,
                    window_size=window_size,
                    shift_size=0 if (i % 2 == 0) else window_size // 2,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    drop=drop,
                    attn_drop=attn_drop,
                    drop_path=drop_path[i]
                    if isinstance(drop_path, list)
                    else drop_path,
                    norm_layer=norm_layer,
                    pretrained_window_size=pretrained_window_size,
                )
                for i in range(depth)
            ]
        )

        # patch merging layer
        if downsample is not None:
            self.downsample = downsample(
                input_resolution, dim=dim, norm_layer=norm_layer
            )
        else:
            self.downsample = None

    def forward(self, x, x_size):
        for blk in self.blocks:
            if self.use_checkpoint:
                x = checkpoint.checkpoint(blk, x, x_size)
            else:
                x = blk(x, x_size)
        if self.downsample is not None:
            x = self.downsample(x)
        return x

    def extra_repr(self) -> str:
        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"

    def flops(self):
        flops = 0
        for blk in self.blocks:
            flops += blk.flops()  # type: ignore
        if self.downsample is not None:
            flops += self.downsample.flops()
        return flops

    def _init_respostnorm(self):
        for blk in self.blocks:
            nn.init.constant_(blk.norm1.bias, 0)  # type: ignore
            nn.init.constant_(blk.norm1.weight, 0)  # type: ignore
            nn.init.constant_(blk.norm2.bias, 0)  # type: ignore
            nn.init.constant_(blk.norm2.weight, 0)  # type: ignore


class PatchEmbed(nn.Module):
    r"""Image to Patch Embedding
    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
        self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]  # type: ignore
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]

        self.in_chans = in_chans
        self.embed_dim = embed_dim

        self.proj = nn.Conv2d(
            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size  # type: ignore
        )
        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        # assert H == self.img_size[0] and W == self.img_size[1],
        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C
        if self.norm is not None:
            x = self.norm(x)
        return x

    def flops(self):
        Ho, Wo = self.patches_resolution
        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])  # type: ignore
        if self.norm is not None:
            flops += Ho * Wo * self.embed_dim
        return flops


class RSTB(nn.Module):
    """Residual Swin Transformer Block (RSTB).

    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
        img_size: Input image size.
        patch_size: Patch size.
        resi_connection: The convolutional block before residual connection.
    """

    def __init__(
        self,
        dim,
        input_resolution,
        depth,
        num_heads,
        window_size,
        mlp_ratio=4.0,
        qkv_bias=True,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        downsample=None,
        use_checkpoint=False,
        img_size=224,
        patch_size=4,
        resi_connection="1conv",
    ):
        super(RSTB, self).__init__()

        self.dim = dim
        self.input_resolution = input_resolution

        self.residual_group = BasicLayer(
            dim=dim,
            input_resolution=input_resolution,
            depth=depth,
            num_heads=num_heads,
            window_size=window_size,
            mlp_ratio=mlp_ratio,
            qkv_bias=qkv_bias,
            drop=drop,
            attn_drop=attn_drop,
            drop_path=drop_path,
            norm_layer=norm_layer,
            downsample=downsample,
            use_checkpoint=use_checkpoint,
        )

        if resi_connection == "1conv":
            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
        elif resi_connection == "3conv":
            # to save parameters and memory
            self.conv = nn.Sequential(
                nn.Conv2d(dim, dim // 4, 3, 1, 1),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(dim // 4, dim, 3, 1, 1),
            )

        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=dim,
            embed_dim=dim,
            norm_layer=None,
        )

        self.patch_unembed = PatchUnEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=dim,
            embed_dim=dim,
            norm_layer=None,
        )

    def forward(self, x, x_size):
        return (
            self.patch_embed(
                self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))
            )
            + x
        )

    def flops(self):
        flops = 0
        flops += self.residual_group.flops()
        H, W = self.input_resolution
        flops += H * W * self.dim * self.dim * 9
        flops += self.patch_embed.flops()
        flops += self.patch_unembed.flops()

        return flops


class PatchUnEmbed(nn.Module):
    r"""Image to Patch Unembedding

    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
        self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]  # type: ignore
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]

        self.in_chans = in_chans
        self.embed_dim = embed_dim

    def forward(self, x, x_size):
        B, HW, C = x.shape
        x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C
        return x

    def flops(self):
        flops = 0
        return flops


class Upsample(nn.Sequential):
    """Upsample module.

    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.
    """

    def __init__(self, scale, num_feat):
        m = []
        if (scale & (scale - 1)) == 0:  # scale = 2^n
            for _ in range(int(math.log(scale, 2))):
                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
                m.append(nn.PixelShuffle(2))
        elif scale == 3:
            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
            m.append(nn.PixelShuffle(3))
        else:
            raise ValueError(
                f"scale {scale} is not supported. " "Supported scales: 2^n and 3."
            )
        super(Upsample, self).__init__(*m)


class Upsample_hf(nn.Sequential):
    """Upsample module.

    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.
    """

    def __init__(self, scale, num_feat):
        m = []
        if (scale & (scale - 1)) == 0:  # scale = 2^n
            for _ in range(int(math.log(scale, 2))):
                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
                m.append(nn.PixelShuffle(2))
        elif scale == 3:
            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
            m.append(nn.PixelShuffle(3))
        else:
            raise ValueError(
                f"scale {scale} is not supported. " "Supported scales: 2^n and 3."
            )
        super(Upsample_hf, self).__init__(*m)


class UpsampleOneStep(nn.Sequential):
    """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
       Used in lightweight SR to save parameters.

    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.

    """

    def __init__(self, scale, num_feat, num_out_ch, input_resolution=None):
        self.num_feat = num_feat
        self.input_resolution = input_resolution
        m = []
        m.append(nn.Conv2d(num_feat, (scale**2) * num_out_ch, 3, 1, 1))
        m.append(nn.PixelShuffle(scale))
        super(UpsampleOneStep, self).__init__(*m)

    def flops(self):
        H, W = self.input_resolution  # type: ignore
        flops = H * W * self.num_feat * 3 * 9
        return flops


class Swin2SR(nn.Module):
    r"""Swin2SR
        A PyTorch impl of : `Swin2SR: SwinV2 Transformer for Compressed Image Super-Resolution and Restoration`.

    Args:
        img_size (int | tuple(int)): Input image size. Default 64
        patch_size (int | tuple(int)): Patch size. Default: 1
        in_chans (int): Number of input image channels. Default: 3
        embed_dim (int): Patch embedding dimension. Default: 96
        depths (tuple(int)): Depth of each Swin Transformer layer.
        num_heads (tuple(int)): Number of attention heads in different layers.
        window_size (int): Window size. Default: 7
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
        drop_rate (float): Dropout rate. Default: 0
        attn_drop_rate (float): Attention dropout rate. Default: 0
        drop_path_rate (float): Stochastic depth rate. Default: 0.1
        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
        patch_norm (bool): If True, add normalization after patch embedding. Default: True
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
        upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
        img_range: Image range. 1. or 255.
        upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
    """

    def __init__(
        self,
        state_dict,
        **kwargs,
    ):
        super(Swin2SR, self).__init__()

        # Defaults
        img_size = 128
        patch_size = 1
        in_chans = 3
        embed_dim = 96
        depths = [6, 6, 6, 6]
        num_heads = [6, 6, 6, 6]
        window_size = 7
        mlp_ratio = 4.0
        qkv_bias = True
        drop_rate = 0.0
        attn_drop_rate = 0.0
        drop_path_rate = 0.1
        norm_layer = nn.LayerNorm
        ape = False
        patch_norm = True
        use_checkpoint = False
        upscale = 2
        img_range = 1.0
        upsampler = ""
        resi_connection = "1conv"
        num_in_ch = in_chans
        num_out_ch = in_chans
        num_feat = 64

        self.model_arch = "Swin2SR"
        self.sub_type = "SR"
        self.state = state_dict
        if "params_ema" in self.state:
            self.state = self.state["params_ema"]
        elif "params" in self.state:
            self.state = self.state["params"]

        state_keys = self.state.keys()

        if "conv_before_upsample.0.weight" in state_keys:
            if "conv_aux.weight" in state_keys:
                upsampler = "pixelshuffle_aux"
            elif "conv_up1.weight" in state_keys:
                upsampler = "nearest+conv"
            else:
                upsampler = "pixelshuffle"
        elif "upsample.0.weight" in state_keys:
            upsampler = "pixelshuffledirect"
        else:
            upsampler = ""

        num_feat = (
            self.state.get("conv_before_upsample.0.weight", None).shape[1]
            if self.state.get("conv_before_upsample.weight", None)
            else 64
        )

        num_in_ch = self.state["conv_first.weight"].shape[1]
        in_chans = num_in_ch
        if "conv_last.weight" in state_keys:
            num_out_ch = self.state["conv_last.weight"].shape[0]
        else:
            num_out_ch = num_in_ch

        upscale = 1
        if upsampler == "nearest+conv":
            upsample_keys = [
                x for x in state_keys if "conv_up" in x and "bias" not in x
            ]

            for upsample_key in upsample_keys:
                upscale *= 2
        elif upsampler == "pixelshuffle" or upsampler == "pixelshuffle_aux":
            upsample_keys = [
                x
                for x in state_keys
                if "upsample" in x and "conv" not in x and "bias" not in x
            ]
            for upsample_key in upsample_keys:
                shape = self.state[upsample_key].shape[0]
                upscale *= math.sqrt(shape // num_feat)
            upscale = int(upscale)
        elif upsampler == "pixelshuffledirect":
            upscale = int(
                math.sqrt(self.state["upsample.0.bias"].shape[0] // num_out_ch)
            )

        max_layer_num = 0
        max_block_num = 0
        for key in state_keys:
            result = re.match(
                r"layers.(\d*).residual_group.blocks.(\d*).norm1.weight", key
            )
            if result:
                layer_num, block_num = result.groups()
                max_layer_num = max(max_layer_num, int(layer_num))
                max_block_num = max(max_block_num, int(block_num))

        depths = [max_block_num + 1 for _ in range(max_layer_num + 1)]

        if (
            "layers.0.residual_group.blocks.0.attn.relative_position_bias_table"
            in state_keys
        ):
            num_heads_num = self.state[
                "layers.0.residual_group.blocks.0.attn.relative_position_bias_table"
            ].shape[-1]
            num_heads = [num_heads_num for _ in range(max_layer_num + 1)]
        else:
            num_heads = depths

        embed_dim = self.state["conv_first.weight"].shape[0]

        mlp_ratio = float(
            self.state["layers.0.residual_group.blocks.0.mlp.fc1.bias"].shape[0]
            / embed_dim
        )

        # TODO: could actually count the layers, but this should do
        if "layers.0.conv.4.weight" in state_keys:
            resi_connection = "3conv"
        else:
            resi_connection = "1conv"

        window_size = int(
            math.sqrt(
                self.state[
                    "layers.0.residual_group.blocks.0.attn.relative_position_index"
                ].shape[0]
            )
        )

        if "layers.0.residual_group.blocks.1.attn_mask" in state_keys:
            img_size = int(
                math.sqrt(
                    self.state["layers.0.residual_group.blocks.1.attn_mask"].shape[0]
                )
                * window_size
            )

        # The JPEG models are the only ones with window-size 7, and they also use this range
        img_range = 255.0 if window_size == 7 else 1.0

        self.in_nc = num_in_ch
        self.out_nc = num_out_ch
        self.num_feat = num_feat
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.depths = depths
        self.window_size = window_size
        self.mlp_ratio = mlp_ratio
        self.scale = upscale
        self.upsampler = upsampler
        self.img_size = img_size
        self.img_range = img_range
        self.resi_connection = resi_connection

        self.supports_fp16 = False  # Too much weirdness to support this at the moment
        self.supports_bfp16 = True
        self.min_size_restriction = 16

        ## END AUTO DETECTION

        if in_chans == 3:
            rgb_mean = (0.4488, 0.4371, 0.4040)
            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
        else:
            self.mean = torch.zeros(1, 1, 1, 1)
        self.upscale = upscale
        self.upsampler = upsampler
        self.window_size = window_size

        #####################################################################################################
        ################################### 1, shallow feature extraction ###################################
        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)

        #####################################################################################################
        ################################### 2, deep feature extraction ######################################
        self.num_layers = len(depths)
        self.embed_dim = embed_dim
        self.ape = ape
        self.patch_norm = patch_norm
        self.num_features = embed_dim
        self.mlp_ratio = mlp_ratio

        # split image into non-overlapping patches
        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=embed_dim,
            embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None,
        )
        num_patches = self.patch_embed.num_patches
        patches_resolution = self.patch_embed.patches_resolution
        self.patches_resolution = patches_resolution

        # merge non-overlapping patches into image
        self.patch_unembed = PatchUnEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=embed_dim,
            embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None,
        )

        # absolute position embedding
        if self.ape:
            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))  # type: ignore
            trunc_normal_(self.absolute_pos_embed, std=0.02)

        self.pos_drop = nn.Dropout(p=drop_rate)

        # stochastic depth
        dpr = [
            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
        ]  # stochastic depth decay rule

        # build Residual Swin Transformer blocks (RSTB)
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):
            layer = RSTB(
                dim=embed_dim,
                input_resolution=(patches_resolution[0], patches_resolution[1]),
                depth=depths[i_layer],
                num_heads=num_heads[i_layer],
                window_size=window_size,
                mlp_ratio=self.mlp_ratio,
                qkv_bias=qkv_bias,
                drop=drop_rate,
                attn_drop=attn_drop_rate,
                drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])],  # type: ignore    # no impact on SR results
                norm_layer=norm_layer,
                downsample=None,
                use_checkpoint=use_checkpoint,
                img_size=img_size,
                patch_size=patch_size,
                resi_connection=resi_connection,
            )
            self.layers.append(layer)

        if self.upsampler == "pixelshuffle_hf":
            self.layers_hf = nn.ModuleList()
            for i_layer in range(self.num_layers):
                layer = RSTB(
                    dim=embed_dim,
                    input_resolution=(patches_resolution[0], patches_resolution[1]),
                    depth=depths[i_layer],
                    num_heads=num_heads[i_layer],
                    window_size=window_size,
                    mlp_ratio=self.mlp_ratio,
                    qkv_bias=qkv_bias,
                    drop=drop_rate,
                    attn_drop=attn_drop_rate,
                    drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])],  # type: ignore    # no impact on SR results # type: ignore
                    norm_layer=norm_layer,
                    downsample=None,
                    use_checkpoint=use_checkpoint,
                    img_size=img_size,
                    patch_size=patch_size,
                    resi_connection=resi_connection,
                )
                self.layers_hf.append(layer)

        self.norm = norm_layer(self.num_features)

        # build the last conv layer in deep feature extraction
        if resi_connection == "1conv":
            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
        elif resi_connection == "3conv":
            # to save parameters and memory
            self.conv_after_body = nn.Sequential(
                nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1),
            )

        #####################################################################################################
        ################################ 3, high quality image reconstruction ################################
        if self.upsampler == "pixelshuffle":
            # for classical SR
            self.conv_before_upsample = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.upsample = Upsample(upscale, num_feat)
            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
        elif self.upsampler == "pixelshuffle_aux":
            self.conv_bicubic = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
            self.conv_before_upsample = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.conv_aux = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
            self.conv_after_aux = nn.Sequential(
                nn.Conv2d(3, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.upsample = Upsample(upscale, num_feat)
            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

        elif self.upsampler == "pixelshuffle_hf":
            self.conv_before_upsample = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.upsample = Upsample(upscale, num_feat)
            self.upsample_hf = Upsample_hf(upscale, num_feat)
            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
            self.conv_first_hf = nn.Sequential(
                nn.Conv2d(num_feat, embed_dim, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.conv_after_body_hf = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
            self.conv_before_upsample_hf = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.conv_last_hf = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

        elif self.upsampler == "pixelshuffledirect":
            # for lightweight SR (to save parameters)
            self.upsample = UpsampleOneStep(
                upscale,
                embed_dim,
                num_out_ch,
                (patches_resolution[0], patches_resolution[1]),
            )
        elif self.upsampler == "nearest+conv":
            # for real-world SR (less artifacts)
            assert self.upscale == 4, "only support x4 now."
            self.conv_before_upsample = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
            self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
            self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
            self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        else:
            # for image denoising and JPEG compression artifact reduction
            self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1)

        self.apply(self._init_weights)

        self.load_state_dict(state_dict)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    @torch.jit.ignore  # type: ignore
    def no_weight_decay(self):
        return {"absolute_pos_embed"}

    @torch.jit.ignore  # type: ignore
    def no_weight_decay_keywords(self):
        return {"relative_position_bias_table"}

    def check_image_size(self, x):
        _, _, h, w = x.size()
        mod_pad_h = (self.window_size - h % self.window_size) % self.window_size
        mod_pad_w = (self.window_size - w % self.window_size) % self.window_size
        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect")
        return x

    def forward_features(self, x):
        x_size = (x.shape[2], x.shape[3])
        x = self.patch_embed(x)
        if self.ape:
            x = x + self.absolute_pos_embed
        x = self.pos_drop(x)

        for layer in self.layers:
            x = layer(x, x_size)

        x = self.norm(x)  # B L C
        x = self.patch_unembed(x, x_size)

        return x

    def forward_features_hf(self, x):
        x_size = (x.shape[2], x.shape[3])
        x = self.patch_embed(x)
        if self.ape:
            x = x + self.absolute_pos_embed
        x = self.pos_drop(x)

        for layer in self.layers_hf:
            x = layer(x, x_size)

        x = self.norm(x)  # B L C
        x = self.patch_unembed(x, x_size)

        return x

    def forward(self, x):
        H, W = x.shape[2:]
        x = self.check_image_size(x)

        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == "pixelshuffle":
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == "pixelshuffle_aux":
            bicubic = F.interpolate(
                x,
                size=(H * self.upscale, W * self.upscale),
                mode="bicubic",
                align_corners=False,
            )
            bicubic = self.conv_bicubic(bicubic)
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            aux = self.conv_aux(x)  # b, 3, LR_H, LR_W
            x = self.conv_after_aux(aux)
            x = (
                self.upsample(x)[:, :, : H * self.upscale, : W * self.upscale]
                + bicubic[:, :, : H * self.upscale, : W * self.upscale]
            )
            x = self.conv_last(x)
            aux = aux / self.img_range + self.mean
        elif self.upsampler == "pixelshuffle_hf":
            # for classical SR with HF
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x_before = self.conv_before_upsample(x)
            x_out = self.conv_last(self.upsample(x_before))

            x_hf = self.conv_first_hf(x_before)
            x_hf = self.conv_after_body_hf(self.forward_features_hf(x_hf)) + x_hf
            x_hf = self.conv_before_upsample_hf(x_hf)
            x_hf = self.conv_last_hf(self.upsample_hf(x_hf))
            x = x_out + x_hf
            x_hf = x_hf / self.img_range + self.mean

        elif self.upsampler == "pixelshuffledirect":
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == "nearest+conv":
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(
                self.conv_up1(
                    torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest")
                )
            )
            x = self.lrelu(
                self.conv_up2(
                    torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest")
                )
            )
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean
        if self.upsampler == "pixelshuffle_aux":
            # NOTE: I removed an "aux" output here. not sure what that was for
            return x[:, :, : H * self.upscale, : W * self.upscale]  # type: ignore

        elif self.upsampler == "pixelshuffle_hf":
            x_out = x_out / self.img_range + self.mean  # type: ignore
            return x_out[:, :, : H * self.upscale, : W * self.upscale], x[:, :, : H * self.upscale, : W * self.upscale], x_hf[:, :, : H * self.upscale, : W * self.upscale]  # type: ignore

        else:
            return x[:, :, : H * self.upscale, : W * self.upscale]

    def flops(self):
        flops = 0
        H, W = self.patches_resolution
        flops += H * W * 3 * self.embed_dim * 9
        flops += self.patch_embed.flops()
        for _i, layer in enumerate(self.layers):
            flops += layer.flops()  # type: ignore
        flops += H * W * 3 * self.embed_dim * self.embed_dim
        flops += self.upsample.flops()  # type: ignore
        return flops

nodes/impl/pytorch/architecture/Swin2SR.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# -----------------------------------------------------------------------------------
# SwinIR: Image Restoration Using Swin Transformer, https://arxiv.org/abs/2108.10257
# Originally Written by Ze Liu, Modified by Jingyun Liang.
# -----------------------------------------------------------------------------------

import math
import re

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint

# Originally from the timm package
from .timm.drop import DropPath
from .timm.helpers import to_2tuple
from .timm.weight_init import trunc_normal_


class Mlp(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        drop=0.0,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


def window_partition(x, window_size):
    """
    Args:
        x: (B, H, W, C)
        window_size (int): window size

    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = (
        x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    )
    return windows


def window_reverse(windows, window_size, H, W):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size (int): Window size
        H (int): Height of image
        W (int): Width of image

    Returns:
        x: (B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(
        B, H // window_size, W // window_size, window_size, window_size, -1
    )
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x


class WindowAttention(nn.Module):
    r"""Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.

    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
    """

    def __init__(
        self,
        dim,
        window_size,
        num_heads,
        qkv_bias=True,
        qk_scale=None,
        attn_drop=0.0,
        proj_drop=0.0,
    ):
        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wh, Ww
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        # define a parameter table of relative position bias
        self.relative_position_bias_table = nn.Parameter(  # type: ignore
            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
        )  # 2*Wh-1 * 2*Ww-1, nH

        # get pair-wise relative position index for each token inside the window
        coords_h = torch.arange(self.window_size[0])
        coords_w = torch.arange(self.window_size[1])
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
        relative_coords = (
            coords_flatten[:, :, None] - coords_flatten[:, None, :]
        )  # 2, Wh*Ww, Wh*Ww
        relative_coords = relative_coords.permute(
            1, 2, 0
        ).contiguous()  # Wh*Ww, Wh*Ww, 2
        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
        self.register_buffer("relative_position_index", relative_position_index)

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)

        self.proj_drop = nn.Dropout(proj_drop)

        trunc_normal_(self.relative_position_bias_table, std=0.02)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x, mask=None):
        """
        Args:
            x: input features with shape of (num_windows*B, N, C)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        B_, N, C = x.shape
        qkv = (
            self.qkv(x)
            .reshape(B_, N, 3, self.num_heads, C // self.num_heads)
            .permute(2, 0, 3, 1, 4)
        )
        q, k, v = (
            qkv[0],
            qkv[1],
            qkv[2],
        )  # make torchscript happy (cannot use tensor as tuple)

        q = q * self.scale
        attn = q @ k.transpose(-2, -1)

        relative_position_bias = self.relative_position_bias_table[
            self.relative_position_index.view(-1)  # type: ignore
        ].view(
            self.window_size[0] * self.window_size[1],
            self.window_size[0] * self.window_size[1],
            -1,
        )  # Wh*Ww,Wh*Ww,nH
        relative_position_bias = relative_position_bias.permute(
            2, 0, 1
        ).contiguous()  # nH, Wh*Ww, Wh*Ww
        attn = attn + relative_position_bias.unsqueeze(0)

        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(
                1
            ).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, N, N)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

    def extra_repr(self) -> str:
        return f"dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}"

    def flops(self, N):
        # calculate flops for 1 window with token length of N
        flops = 0
        # qkv = self.qkv(x)
        flops += N * self.dim * 3 * self.dim
        # attn = (q @ k.transpose(-2, -1))
        flops += self.num_heads * N * (self.dim // self.num_heads) * N
        #  x = (attn @ v)
        flops += self.num_heads * N * N * (self.dim // self.num_heads)
        # x = self.proj(x)
        flops += N * self.dim * self.dim
        return flops


class SwinTransformerBlock(nn.Module):
    r"""Swin Transformer Block.

    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resulotion.
        num_heads (int): Number of attention heads.
        window_size (int): Window size.
        shift_size (int): Shift size for SW-MSA.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float, optional): Stochastic depth rate. Default: 0.0
        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(
        self,
        dim,
        input_resolution,
        num_heads,
        window_size=7,
        shift_size=0,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        if min(self.input_resolution) <= self.window_size:
            # if window size is larger than input resolution, we don't partition windows
            self.shift_size = 0
            self.window_size = min(self.input_resolution)
        assert (
            0 <= self.shift_size < self.window_size
        ), "shift_size must in 0-window_size"

        self.norm1 = norm_layer(dim)
        self.attn = WindowAttention(
            dim,
            window_size=to_2tuple(self.window_size),
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            attn_drop=attn_drop,
            proj_drop=drop,
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

        if self.shift_size > 0:
            attn_mask = self.calculate_mask(self.input_resolution)
        else:
            attn_mask = None

        self.register_buffer("attn_mask", attn_mask)

    def calculate_mask(self, x_size):
        # calculate attention mask for SW-MSA
        H, W = x_size
        img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
        h_slices = (
            slice(0, -self.window_size),
            slice(-self.window_size, -self.shift_size),
            slice(-self.shift_size, None),
        )
        w_slices = (
            slice(0, -self.window_size),
            slice(-self.window_size, -self.shift_size),
            slice(-self.shift_size, None),
        )
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                img_mask[:, h, w, :] = cnt
                cnt += 1

        mask_windows = window_partition(
            img_mask, self.window_size
        )  # nW, window_size, window_size, 1
        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
            attn_mask == 0, float(0.0)
        )

        return attn_mask

    def forward(self, x, x_size):
        H, W = x_size
        B, L, C = x.shape
        # assert L == H * W, "input feature has wrong size"

        shortcut = x
        x = self.norm1(x)
        x = x.view(B, H, W, C)

        # cyclic shift
        if self.shift_size > 0:
            shifted_x = torch.roll(
                x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)
            )
        else:
            shifted_x = x

        # partition windows
        x_windows = window_partition(
            shifted_x, self.window_size
        )  # nW*B, window_size, window_size, C
        x_windows = x_windows.view(
            -1, self.window_size * self.window_size, C
        )  # nW*B, window_size*window_size, C

        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
        if self.input_resolution == x_size:
            attn_windows = self.attn(
                x_windows, mask=self.attn_mask
            )  # nW*B, window_size*window_size, C
        else:
            attn_windows = self.attn(
                x_windows, mask=self.calculate_mask(x_size).to(x.device)
            )

        # merge windows
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C

        # reverse cyclic shift
        if self.shift_size > 0:
            x = torch.roll(
                shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)
            )
        else:
            x = shifted_x
        x = x.view(B, H * W, C)

        # FFN
        x = shortcut + self.drop_path(x)
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x

    def extra_repr(self) -> str:
        return (
            f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, "
            f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
        )

    def flops(self):
        flops = 0
        H, W = self.input_resolution
        # norm1
        flops += self.dim * H * W
        # W-MSA/SW-MSA
        nW = H * W / self.window_size / self.window_size
        flops += nW * self.attn.flops(self.window_size * self.window_size)
        # mlp
        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
        # norm2
        flops += self.dim * H * W
        return flops


class PatchMerging(nn.Module):
    r"""Patch Merging Layer.

    Args:
        input_resolution (tuple[int]): Resolution of input feature.
        dim (int): Number of input channels.
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(4 * dim)

    def forward(self, x):
        """
        x: B, H*W, C
        """
        H, W = self.input_resolution
        B, L, C = x.shape
        assert L == H * W, "input feature has wrong size"
        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."

        x = x.view(B, H, W, C)

        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C

        x = self.norm(x)
        x = self.reduction(x)

        return x

    def extra_repr(self) -> str:
        return f"input_resolution={self.input_resolution}, dim={self.dim}"

    def flops(self):
        H, W = self.input_resolution
        flops = H * W * self.dim
        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
        return flops


class BasicLayer(nn.Module):
    """A basic Swin Transformer layer for one stage.

    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
    """

    def __init__(
        self,
        dim,
        input_resolution,
        depth,
        num_heads,
        window_size,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        downsample=None,
        use_checkpoint=False,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth
        self.use_checkpoint = use_checkpoint

        # build blocks
        self.blocks = nn.ModuleList(
            [
                SwinTransformerBlock(
                    dim=dim,
                    input_resolution=input_resolution,
                    num_heads=num_heads,
                    window_size=window_size,
                    shift_size=0 if (i % 2 == 0) else window_size // 2,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop,
                    attn_drop=attn_drop,
                    drop_path=drop_path[i]
                    if isinstance(drop_path, list)
                    else drop_path,
                    norm_layer=norm_layer,
                )
                for i in range(depth)
            ]
        )

        # patch merging layer
        if downsample is not None:
            self.downsample = downsample(
                input_resolution, dim=dim, norm_layer=norm_layer
            )
        else:
            self.downsample = None

    def forward(self, x, x_size):
        for blk in self.blocks:
            if self.use_checkpoint:
                x = checkpoint.checkpoint(blk, x, x_size)
            else:
                x = blk(x, x_size)
        if self.downsample is not None:
            x = self.downsample(x)
        return x

    def extra_repr(self) -> str:
        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"

    def flops(self):
        flops = 0
        for blk in self.blocks:
            flops += blk.flops()  # type: ignore
        if self.downsample is not None:
            flops += self.downsample.flops()
        return flops


class RSTB(nn.Module):
    """Residual Swin Transformer Block (RSTB).

    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
        img_size: Input image size.
        patch_size: Patch size.
        resi_connection: The convolutional block before residual connection.
    """

    def __init__(
        self,
        dim,
        input_resolution,
        depth,
        num_heads,
        window_size,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        downsample=None,
        use_checkpoint=False,
        img_size=224,
        patch_size=4,
        resi_connection="1conv",
    ):
        super(RSTB, self).__init__()

        self.dim = dim
        self.input_resolution = input_resolution

        self.residual_group = BasicLayer(
            dim=dim,
            input_resolution=input_resolution,
            depth=depth,
            num_heads=num_heads,
            window_size=window_size,
            mlp_ratio=mlp_ratio,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            drop=drop,
            attn_drop=attn_drop,
            drop_path=drop_path,
            norm_layer=norm_layer,
            downsample=downsample,
            use_checkpoint=use_checkpoint,
        )

        if resi_connection == "1conv":
            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
        elif resi_connection == "3conv":
            # to save parameters and memory
            self.conv = nn.Sequential(
                nn.Conv2d(dim, dim // 4, 3, 1, 1),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(dim // 4, dim, 3, 1, 1),
            )

        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=0,
            embed_dim=dim,
            norm_layer=None,
        )

        self.patch_unembed = PatchUnEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=0,
            embed_dim=dim,
            norm_layer=None,
        )

    def forward(self, x, x_size):
        return (
            self.patch_embed(
                self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))
            )
            + x
        )

    def flops(self):
        flops = 0
        flops += self.residual_group.flops()
        H, W = self.input_resolution
        flops += H * W * self.dim * self.dim * 9
        flops += self.patch_embed.flops()
        flops += self.patch_unembed.flops()

        return flops


class PatchEmbed(nn.Module):
    r"""Image to Patch Embedding

    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
        self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [
            img_size[0] // patch_size[0],  # type: ignore
            img_size[1] // patch_size[1],  # type: ignore
        ]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]

        self.in_chans = in_chans
        self.embed_dim = embed_dim

        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        x = x.flatten(2).transpose(1, 2)  # B Ph*Pw C
        if self.norm is not None:
            x = self.norm(x)
        return x

    def flops(self):
        flops = 0
        H, W = self.img_size
        if self.norm is not None:
            flops += H * W * self.embed_dim  # type: ignore
        return flops


class PatchUnEmbed(nn.Module):
    r"""Image to Patch Unembedding

    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
        self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [
            img_size[0] // patch_size[0],  # type: ignore
            img_size[1] // patch_size[1],  # type: ignore
        ]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]

        self.in_chans = in_chans
        self.embed_dim = embed_dim

    def forward(self, x, x_size):
        B, HW, C = x.shape
        x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C
        return x

    def flops(self):
        flops = 0
        return flops


class Upsample(nn.Sequential):
    """Upsample module.

    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.
    """

    def __init__(self, scale, num_feat):
        m = []
        if (scale & (scale - 1)) == 0:  # scale = 2^n
            for _ in range(int(math.log(scale, 2))):
                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
                m.append(nn.PixelShuffle(2))
        elif scale == 3:
            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
            m.append(nn.PixelShuffle(3))
        else:
            raise ValueError(
                f"scale {scale} is not supported. " "Supported scales: 2^n and 3."
            )
        super(Upsample, self).__init__(*m)


class UpsampleOneStep(nn.Sequential):
    """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
       Used in lightweight SR to save parameters.

    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.

    """

    def __init__(self, scale, num_feat, num_out_ch, input_resolution=None):
        self.num_feat = num_feat
        self.input_resolution = input_resolution
        m = []
        m.append(nn.Conv2d(num_feat, (scale**2) * num_out_ch, 3, 1, 1))
        m.append(nn.PixelShuffle(scale))
        super(UpsampleOneStep, self).__init__(*m)

    def flops(self):
        H, W = self.input_resolution  # type: ignore
        flops = H * W * self.num_feat * 3 * 9
        return flops


class SwinIR(nn.Module):
    r"""SwinIR
        A PyTorch impl of : `SwinIR: Image Restoration Using Swin Transformer`, based on Swin Transformer.

    Args:
        img_size (int | tuple(int)): Input image size. Default 64
        patch_size (int | tuple(int)): Patch size. Default: 1
        in_chans (int): Number of input image channels. Default: 3
        embed_dim (int): Patch embedding dimension. Default: 96
        depths (tuple(int)): Depth of each Swin Transformer layer.
        num_heads (tuple(int)): Number of attention heads in different layers.
        window_size (int): Window size. Default: 7
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
        drop_rate (float): Dropout rate. Default: 0
        attn_drop_rate (float): Attention dropout rate. Default: 0
        drop_path_rate (float): Stochastic depth rate. Default: 0.1
        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
        patch_norm (bool): If True, add normalization after patch embedding. Default: True
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
        upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
        img_range: Image range. 1. or 255.
        upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
    """

    def __init__(
        self,
        state_dict,
        **kwargs,
    ):
        super(SwinIR, self).__init__()

        # Defaults
        img_size = 64
        patch_size = 1
        in_chans = 3
        embed_dim = 96
        depths = [6, 6, 6, 6]
        num_heads = [6, 6, 6, 6]
        window_size = 7
        mlp_ratio = 4.0
        qkv_bias = True
        qk_scale = None
        drop_rate = 0.0
        attn_drop_rate = 0.0
        drop_path_rate = 0.1
        norm_layer = nn.LayerNorm
        ape = False
        patch_norm = True
        use_checkpoint = False
        upscale = 2
        img_range = 1.0
        upsampler = ""
        resi_connection = "1conv"
        num_feat = 64
        num_in_ch = in_chans
        num_out_ch = in_chans
        self.start_unshuffle = 1

        self.model_arch = "SwinIR"
        self.sub_type = "SR"
        self.state = state_dict
        if "params_ema" in self.state:
            self.state = self.state["params_ema"]
        elif "params" in self.state:
            self.state = self.state["params"]

        state_keys = self.state.keys()

        if "conv_before_upsample.0.weight" in state_keys:
            if "conv_up1.weight" in state_keys:
                upsampler = "nearest+conv"
            else:
                upsampler = "pixelshuffle"
        elif "upsample.0.weight" in state_keys:
            upsampler = "pixelshuffledirect"
        else:
            upsampler = ""

        num_feat = (
            self.state.get("conv_before_upsample.0.weight", None).shape[1]
            if self.state.get("conv_before_upsample.weight", None)
            else 64
        )

        if "conv_first.1.weight" in self.state:
            self.state["conv_first.weight"] = self.state.pop("conv_first.1.weight")
            self.state["conv_first.bias"] = self.state.pop("conv_first.1.bias")
            self.start_unshuffle = round(
                math.sqrt(self.state["conv_first.weight"].shape[1] // 3)
            )

        num_in_ch = self.state["conv_first.weight"].shape[1]
        in_chans = num_in_ch
        if "conv_last.weight" in state_keys:
            num_out_ch = self.state["conv_last.weight"].shape[0]
        else:
            num_out_ch = num_in_ch

        upscale = 1
        if upsampler == "nearest+conv":
            upsample_keys = [
                x for x in state_keys if "conv_up" in x and "bias" not in x
            ]

            for upsample_key in upsample_keys:
                upscale *= 2
        elif upsampler == "pixelshuffle":
            upsample_keys = [
                x
                for x in state_keys
                if "upsample" in x and "conv" not in x and "bias" not in x
            ]
            for upsample_key in upsample_keys:
                shape = self.state[upsample_key].shape[0]
                upscale *= math.sqrt(shape // num_feat)
            upscale = int(upscale)
        elif upsampler == "pixelshuffledirect":
            upscale = int(
                math.sqrt(self.state["upsample.0.bias"].shape[0] // num_out_ch)
            )

        max_layer_num = 0
        max_block_num = 0
        for key in state_keys:
            result = re.match(
                r"layers.(\d*).residual_group.blocks.(\d*).norm1.weight", key
            )
            if result:
                layer_num, block_num = result.groups()
                max_layer_num = max(max_layer_num, int(layer_num))
                max_block_num = max(max_block_num, int(block_num))

        depths = [max_block_num + 1 for _ in range(max_layer_num + 1)]

        if (
            "layers.0.residual_group.blocks.0.attn.relative_position_bias_table"
            in state_keys
        ):
            num_heads_num = self.state[
                "layers.0.residual_group.blocks.0.attn.relative_position_bias_table"
            ].shape[-1]
            num_heads = [num_heads_num for _ in range(max_layer_num + 1)]
        else:
            num_heads = depths

        embed_dim = self.state["conv_first.weight"].shape[0]

        mlp_ratio = float(
            self.state["layers.0.residual_group.blocks.0.mlp.fc1.bias"].shape[0]
            / embed_dim
        )

        # TODO: could actually count the layers, but this should do
        if "layers.0.conv.4.weight" in state_keys:
            resi_connection = "3conv"
        else:
            resi_connection = "1conv"

        window_size = int(
            math.sqrt(
                self.state[
                    "layers.0.residual_group.blocks.0.attn.relative_position_index"
                ].shape[0]
            )
        )

        if "layers.0.residual_group.blocks.1.attn_mask" in state_keys:
            img_size = int(
                math.sqrt(
                    self.state["layers.0.residual_group.blocks.1.attn_mask"].shape[0]
                )
                * window_size
            )

        # The JPEG models are the only ones with window-size 7, and they also use this range
        img_range = 255.0 if window_size == 7 else 1.0

        self.in_nc = num_in_ch // self.start_unshuffle**2
        self.out_nc = num_out_ch
        self.num_feat = num_feat
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.depths = depths
        self.window_size = window_size
        self.mlp_ratio = mlp_ratio
        self.scale = upscale
        self.upsampler = upsampler
        self.img_size = img_size
        self.img_range = img_range
        self.resi_connection = resi_connection

        self.supports_fp16 = False  # Too much weirdness to support this at the moment
        self.supports_bfp16 = True
        self.min_size_restriction = 16

        self.img_range = img_range
        if in_chans == 3:
            rgb_mean = (0.4488, 0.4371, 0.4040)
            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
        else:
            self.mean = torch.zeros(1, 1, 1, 1)
        self.upscale = upscale
        self.upsampler = upsampler
        self.window_size = window_size

        #####################################################################################################
        ################################### 1, shallow feature extraction ###################################
        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)

        #####################################################################################################
        ################################### 2, deep feature extraction ######################################
        self.num_layers = len(depths)
        self.embed_dim = embed_dim
        self.ape = ape
        self.patch_norm = patch_norm
        self.num_features = embed_dim
        self.mlp_ratio = mlp_ratio

        # split image into non-overlapping patches
        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=embed_dim,
            embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None,
        )
        num_patches = self.patch_embed.num_patches
        patches_resolution = self.patch_embed.patches_resolution
        self.patches_resolution = patches_resolution

        # merge non-overlapping patches into image
        self.patch_unembed = PatchUnEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=embed_dim,
            embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None,
        )

        # absolute position embedding
        if self.ape:
            self.absolute_pos_embed = nn.Parameter(  # type: ignore
                torch.zeros(1, num_patches, embed_dim)
            )
            trunc_normal_(self.absolute_pos_embed, std=0.02)

        self.pos_drop = nn.Dropout(p=drop_rate)

        # stochastic depth
        dpr = [
            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
        ]  # stochastic depth decay rule

        # build Residual Swin Transformer blocks (RSTB)
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):
            layer = RSTB(
                dim=embed_dim,
                input_resolution=(patches_resolution[0], patches_resolution[1]),
                depth=depths[i_layer],
                num_heads=num_heads[i_layer],
                window_size=window_size,
                mlp_ratio=self.mlp_ratio,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop=drop_rate,
                attn_drop=attn_drop_rate,
                drop_path=dpr[
                    sum(depths[:i_layer]) : sum(depths[: i_layer + 1])  # type: ignore
                ],  # no impact on SR results
                norm_layer=norm_layer,
                downsample=None,
                use_checkpoint=use_checkpoint,
                img_size=img_size,
                patch_size=patch_size,
                resi_connection=resi_connection,
            )
            self.layers.append(layer)
        self.norm = norm_layer(self.num_features)

        # build the last conv layer in deep feature extraction
        if resi_connection == "1conv":
            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
        elif resi_connection == "3conv":
            # to save parameters and memory
            self.conv_after_body = nn.Sequential(
                nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),
                nn.LeakyReLU(negative_slope=0.2, inplace=True),
                nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1),
            )

        #####################################################################################################
        ################################ 3, high quality image reconstruction ################################
        if self.upsampler == "pixelshuffle":
            # for classical SR
            self.conv_before_upsample = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.upsample = Upsample(upscale, num_feat)
            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
        elif self.upsampler == "pixelshuffledirect":
            # for lightweight SR (to save parameters)
            self.upsample = UpsampleOneStep(
                upscale,
                embed_dim,
                num_out_ch,
                (patches_resolution[0], patches_resolution[1]),
            )
        elif self.upsampler == "nearest+conv":
            # for real-world SR (less artifacts)
            self.conv_before_upsample = nn.Sequential(
                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
            )
            self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
            if self.upscale == 4:
                self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
            elif self.upscale == 8:
                self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
                self.conv_up3 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
            self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
            self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        else:
            # for image denoising and JPEG compression artifact reduction
            self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1)

        self.apply(self._init_weights)
        self.load_state_dict(self.state, strict=False)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    @torch.jit.ignore  # type: ignore
    def no_weight_decay(self):
        return {"absolute_pos_embed"}

    @torch.jit.ignore  # type: ignore
    def no_weight_decay_keywords(self):
        return {"relative_position_bias_table"}

    def check_image_size(self, x):
        _, _, h, w = x.size()
        mod_pad_h = (self.window_size - h % self.window_size) % self.window_size
        mod_pad_w = (self.window_size - w % self.window_size) % self.window_size
        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect")
        return x

    def forward_features(self, x):
        x_size = (x.shape[2], x.shape[3])
        x = self.patch_embed(x)
        if self.ape:
            x = x + self.absolute_pos_embed
        x = self.pos_drop(x)

        for layer in self.layers:
            x = layer(x, x_size)

        x = self.norm(x)  # B L C
        x = self.patch_unembed(x, x_size)

        return x

    def forward(self, x):
        H, W = x.shape[2:]
        x = self.check_image_size(x)

        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.start_unshuffle > 1:
            up = torch.nn.Upsample(scale_factor=self.start_unshuffle, mode="bicubic")
            x = up(x)
            x = torch.nn.functional.pixel_unshuffle(x, self.start_unshuffle)

        if self.upsampler == "pixelshuffle":
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == "pixelshuffledirect":
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == "nearest+conv":
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(
                self.conv_up1(
                    torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest")  # type: ignore
                )
            )
            if self.upscale == 4:
                x = self.lrelu(
                    self.conv_up2(
                        torch.nn.functional.interpolate(  # type: ignore
                            x, scale_factor=2, mode="nearest"
                        )
                    )
                )
            elif self.upscale == 8:
                x = self.lrelu(
                    self.conv_up2(
                        torch.nn.functional.interpolate(
                            x, scale_factor=2, mode="nearest"
                        )
                    )
                )
                x = self.lrelu(
                    self.conv_up3(
                        torch.nn.functional.interpolate(
                            x, scale_factor=2, mode="nearest"
                        )
                    )
                )
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean

        return x[:, :, : H * self.upscale, : W * self.upscale]

    def flops(self):
        flops = 0
        H, W = self.patches_resolution
        flops += H * W * 3 * self.embed_dim * 9
        flops += self.patch_embed.flops()
        for _i, layer in enumerate(self.layers):
            flops += layer.flops()  # type: ignore
        flops += H * W * 3 * self.embed_dim * self.embed_dim
        flops += self.upsample.flops()  # type: ignore
        return flops

nodes/impl/pytorch/architecture/SwinIR.py
vladmandic/sd-extension-chainner
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import annotations

from collections import OrderedDict
from typing import Literal

import torch
import torch.nn as nn

####################
# Basic blocks
####################


def act(act_type: str, inplace=True, neg_slope=0.2, n_prelu=1):
    # helper selecting activation
    # neg_slope: for leakyrelu and init of prelu
    # n_prelu: for p_relu num_parameters
    act_type = act_type.lower()
    if act_type == "relu":
        layer = nn.ReLU(inplace)
    elif act_type == "leakyrelu":
        layer = nn.LeakyReLU(neg_slope, inplace)
    elif act_type == "prelu":
        layer = nn.PReLU(num_parameters=n_prelu, init=neg_slope)
    else:
        raise NotImplementedError(
            "activation layer [{:s}] is not found".format(act_type)
        )
    return layer


def norm(norm_type: str, nc: int):
    # helper selecting normalization layer
    norm_type = norm_type.lower()
    if norm_type == "batch":
        layer = nn.BatchNorm2d(nc, affine=True)
    elif norm_type == "instance":
        layer = nn.InstanceNorm2d(nc, affine=False)
    else:
        raise NotImplementedError(
            "normalization layer [{:s}] is not found".format(norm_type)
        )
    return layer


def pad(pad_type: str, padding):
    # helper selecting padding layer
    # if padding is 'zero', do by conv layers
    pad_type = pad_type.lower()
    if padding == 0:
        return None
    if pad_type == "reflect":
        layer = nn.ReflectionPad2d(padding)
    elif pad_type == "replicate":
        layer = nn.ReplicationPad2d(padding)
    else:
        raise NotImplementedError(
            "padding layer [{:s}] is not implemented".format(pad_type)
        )
    return layer


def get_valid_padding(kernel_size, dilation):
    kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1)
    padding = (kernel_size - 1) // 2
    return padding


class ConcatBlock(nn.Module):
    # Concat the output of a submodule to its input
    def __init__(self, submodule):
        super(ConcatBlock, self).__init__()
        self.sub = submodule

    def forward(self, x):
        output = torch.cat((x, self.sub(x)), dim=1)
        return output

    def __repr__(self):
        tmpstr = "Identity .. \n|"
        modstr = self.sub.__repr__().replace("\n", "\n|")
        tmpstr = tmpstr + modstr
        return tmpstr


class ShortcutBlock(nn.Module):
    # Elementwise sum the output of a submodule to its input
    def __init__(self, submodule):
        super(ShortcutBlock, self).__init__()
        self.sub = submodule

    def forward(self, x):
        output = x + self.sub(x)
        return output

    def __repr__(self):
        tmpstr = "Identity + \n|"
        modstr = self.sub.__repr__().replace("\n", "\n|")
        tmpstr = tmpstr + modstr
        return tmpstr


class ShortcutBlockSPSR(nn.Module):
    # Elementwise sum the output of a submodule to its input
    def __init__(self, submodule):
        super(ShortcutBlockSPSR, self).__init__()
        self.sub = submodule

    def forward(self, x):
        return x, self.sub

    def __repr__(self):
        tmpstr = "Identity + \n|"
        modstr = self.sub.__repr__().replace("\n", "\n|")
        tmpstr = tmpstr + modstr
        return tmpstr


def sequential(*args):
    # Flatten Sequential. It unwraps nn.Sequential.
    if len(args) == 1:
        if isinstance(args[0], OrderedDict):
            raise NotImplementedError("sequential does not support OrderedDict input.")
        return args[0]  # No sequential is needed.
    modules = []
    for module in args:
        if isinstance(module, nn.Sequential):
            for submodule in module.children():
                modules.append(submodule)
        elif isinstance(module, nn.Module):
            modules.append(module)
    return nn.Sequential(*modules)


ConvMode = Literal["CNA", "NAC", "CNAC"]


# 2x2x2 Conv Block
def conv_block_2c2(
    in_nc,
    out_nc,
    act_type="relu",
):
    return sequential(
        nn.Conv2d(in_nc, out_nc, kernel_size=2, padding=1),
        nn.Conv2d(out_nc, out_nc, kernel_size=2, padding=0),
        act(act_type) if act_type else None,
    )


def conv_block(
    in_nc: int,
    out_nc: int,
    kernel_size,
    stride=1,
    dilation=1,
    groups=1,
    bias=True,
    pad_type="zero",
    norm_type: str | None = None,
    act_type: str | None = "relu",
    mode: ConvMode = "CNA",
    c2x2=False,
):
    """
    Conv layer with padding, normalization, activation
    mode: CNA --> Conv -> Norm -> Act
        NAC --> Norm -> Act --> Conv (Identity Mappings in Deep Residual Networks, ECCV16)
    """

    if c2x2:
        return conv_block_2c2(in_nc, out_nc, act_type=act_type)

    assert mode in ("CNA", "NAC", "CNAC"), "Wrong conv mode [{:s}]".format(mode)
    padding = get_valid_padding(kernel_size, dilation)
    p = pad(pad_type, padding) if pad_type and pad_type != "zero" else None
    padding = padding if pad_type == "zero" else 0

    c = nn.Conv2d(
        in_nc,
        out_nc,
        kernel_size=kernel_size,
        stride=stride,
        padding=padding,
        dilation=dilation,
        bias=bias,
        groups=groups,
    )
    a = act(act_type) if act_type else None
    if mode in ("CNA", "CNAC"):
        n = norm(norm_type, out_nc) if norm_type else None
        return sequential(p, c, n, a)
    elif mode == "NAC":
        if norm_type is None and act_type is not None:
            a = act(act_type, inplace=False)
            # Important!
            # input----ReLU(inplace)----Conv--+----output
            #        |________________________|
            # inplace ReLU will modify the input, therefore wrong output
        n = norm(norm_type, in_nc) if norm_type else None
        return sequential(n, a, p, c)
    else:
        raise AssertionError(f"Invalid conv mode {mode}")


####################
# Useful blocks
####################


class ResNetBlock(nn.Module):
    """
    ResNet Block, 3-3 style
    with extra residual scaling used in EDSR
    (Enhanced Deep Residual Networks for Single Image Super-Resolution, CVPRW 17)
    """

    def __init__(
        self,
        in_nc,
        mid_nc,
        out_nc,
        kernel_size=3,
        stride=1,
        dilation=1,
        groups=1,
        bias=True,
        pad_type="zero",
        norm_type=None,
        act_type="relu",
        mode: ConvMode = "CNA",
        res_scale=1,
    ):
        super(ResNetBlock, self).__init__()
        conv0 = conv_block(
            in_nc,
            mid_nc,
            kernel_size,
            stride,
            dilation,
            groups,
            bias,
            pad_type,
            norm_type,
            act_type,
            mode,
        )
        if mode == "CNA":
            act_type = None
        if mode == "CNAC":  # Residual path: |-CNAC-|
            act_type = None
            norm_type = None
        conv1 = conv_block(
            mid_nc,
            out_nc,
            kernel_size,
            stride,
            dilation,
            groups,
            bias,
            pad_type,
            norm_type,
            act_type,
            mode,
        )
        # if in_nc != out_nc:
        #     self.project = conv_block(in_nc, out_nc, 1, stride, dilation, 1, bias, pad_type, \
        #         None, None)
        #     print('Need a projecter in ResNetBlock.')
        # else:
        #     self.project = lambda x:x
        self.res = sequential(conv0, conv1)
        self.res_scale = res_scale

    def forward(self, x):
        res = self.res(x).mul(self.res_scale)
        return x + res


class RRDB(nn.Module):
    """
    Residual in Residual Dense Block
    (ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks)
    """

    def __init__(
        self,
        nf,
        kernel_size=3,
        gc=32,
        stride=1,
        bias: bool = True,
        pad_type="zero",
        norm_type=None,
        act_type="leakyrelu",
        mode: ConvMode = "CNA",
        _convtype="Conv2D",
        _spectral_norm=False,
        plus=False,
        c2x2=False,
    ):
        super(RRDB, self).__init__()
        self.RDB1 = ResidualDenseBlock_5C(
            nf,
            kernel_size,
            gc,
            stride,
            bias,
            pad_type,
            norm_type,
            act_type,
            mode,
            plus=plus,
            c2x2=c2x2,
        )
        self.RDB2 = ResidualDenseBlock_5C(
            nf,
            kernel_size,
            gc,
            stride,
            bias,
            pad_type,
            norm_type,
            act_type,
            mode,
            plus=plus,
            c2x2=c2x2,
        )
        self.RDB3 = ResidualDenseBlock_5C(
            nf,
            kernel_size,
            gc,
            stride,
            bias,
            pad_type,
            norm_type,
            act_type,
            mode,
            plus=plus,
            c2x2=c2x2,
        )

    def forward(self, x):
        out = self.RDB1(x)
        out = self.RDB2(out)
        out = self.RDB3(out)
        return out * 0.2 + x


class ResidualDenseBlock_5C(nn.Module):
    """
    Residual Dense Block
    style: 5 convs
    The core module of paper: (Residual Dense Network for Image Super-Resolution, CVPR 18)
    Modified options that can be used:
        - "Partial Convolution based Padding" arXiv:1811.11718
        - "Spectral normalization" arXiv:1802.05957
        - "ICASSP 2020 - ESRGAN+ : Further Improving ESRGAN" N. C.
            {Rakotonirina} and A. {Rasoanaivo}

    Args:
        nf (int): Channel number of intermediate features (num_feat).
        gc (int): Channels for each growth (num_grow_ch: growth channel,
            i.e. intermediate channels).
        convtype (str): the type of convolution to use. Default: 'Conv2D'
        gaussian_noise (bool): enable the ESRGAN+ gaussian noise (no new
            trainable parameters)
        plus (bool): enable the additional residual paths from ESRGAN+
            (adds trainable parameters)
    """

    def __init__(
        self,
        nf=64,
        kernel_size=3,
        gc=32,
        stride=1,
        bias: bool = True,
        pad_type="zero",
        norm_type=None,
        act_type="leakyrelu",
        mode: ConvMode = "CNA",
        plus=False,
        c2x2=False,
    ):
        super(ResidualDenseBlock_5C, self).__init__()

        ## +
        self.conv1x1 = conv1x1(nf, gc) if plus else None
        ## +

        self.conv1 = conv_block(
            nf,
            gc,
            kernel_size,
            stride,
            bias=bias,
            pad_type=pad_type,
            norm_type=norm_type,
            act_type=act_type,
            mode=mode,
            c2x2=c2x2,
        )
        self.conv2 = conv_block(
            nf + gc,
            gc,
            kernel_size,
            stride,
            bias=bias,
            pad_type=pad_type,
            norm_type=norm_type,
            act_type=act_type,
            mode=mode,
            c2x2=c2x2,
        )
        self.conv3 = conv_block(
            nf + 2 * gc,
            gc,
            kernel_size,
            stride,
            bias=bias,
            pad_type=pad_type,
            norm_type=norm_type,
            act_type=act_type,
            mode=mode,
            c2x2=c2x2,
        )
        self.conv4 = conv_block(
            nf + 3 * gc,
            gc,
            kernel_size,
            stride,
            bias=bias,
            pad_type=pad_type,
            norm_type=norm_type,
            act_type=act_type,
            mode=mode,
            c2x2=c2x2,
        )
        if mode == "CNA":
            last_act = None
        else:
            last_act = act_type
        self.conv5 = conv_block(
            nf + 4 * gc,
            nf,
            3,
            stride,
            bias=bias,
            pad_type=pad_type,
            norm_type=norm_type,
            act_type=last_act,
            mode=mode,
            c2x2=c2x2,
        )

    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(torch.cat((x, x1), 1))
        if self.conv1x1:
            # pylint: disable=not-callable
            x2 = x2 + self.conv1x1(x)  # +
        x3 = self.conv3(torch.cat((x, x1, x2), 1))
        x4 = self.conv4(torch.cat((x, x1, x2, x3), 1))
        if self.conv1x1:
            x4 = x4 + x2  # +
        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
        return x5 * 0.2 + x


def conv1x1(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


####################
# Upsampler
####################


def pixelshuffle_block(
    in_nc: int,
    out_nc: int,
    upscale_factor=2,
    kernel_size=3,
    stride=1,
    bias=True,
    pad_type="zero",
    norm_type: str | None = None,
    act_type="relu",
):
    """
    Pixel shuffle layer
    (Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional
    Neural Network, CVPR17)
    """
    conv = conv_block(
        in_nc,
        out_nc * (upscale_factor**2),
        kernel_size,
        stride,
        bias=bias,
        pad_type=pad_type,
        norm_type=None,
        act_type=None,
    )
    pixel_shuffle = nn.PixelShuffle(upscale_factor)

    n = norm(norm_type, out_nc) if norm_type else None
    a = act(act_type) if act_type else None
    return sequential(conv, pixel_shuffle, n, a)


def upconv_block(
    in_nc: int,
    out_nc: int,
    upscale_factor=2,
    kernel_size=3,
    stride=1,
    bias=True,
    pad_type="zero",
    norm_type: str | None = None,
    act_type="relu",
    mode="nearest",
    c2x2=False,
):
    # Up conv
    # described in https://distill.pub/2016/deconv-checkerboard/
    upsample = nn.Upsample(scale_factor=upscale_factor, mode=mode)
    conv = conv_block(
        in_nc,
        out_nc,
        kernel_size,
        stride,
        bias=bias,
        pad_type=pad_type,
        norm_type=norm_type,
        act_type=act_type,
        c2x2=c2x2,
    )
    return sequential(upsample, conv)

nodes/impl/pytorch/architecture/block.py
vladmandic/sd-extension-chainner
import torch.nn as nn


def conv3x3(inplanes, outplanes, stride=1):
    """A simple wrapper for 3x3 convolution with padding.

    Args:
        inplanes (int): Channel number of inputs.
        outplanes (int): Channel number of outputs.
        stride (int): Stride in convolution. Default: 1.
    """
    return nn.Conv2d(
        inplanes, outplanes, kernel_size=3, stride=stride, padding=1, bias=False
    )


class BasicBlock(nn.Module):
    """Basic residual block used in the ResNetArcFace architecture.

    Args:
        inplanes (int): Channel number of inputs.
        planes (int): Channel number of outputs.
        stride (int): Stride in convolution. Default: 1.
        downsample (nn.Module): The downsample module. Default: None.
    """

    expansion = 1  # output channel expansion ratio

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class IRBlock(nn.Module):
    """Improved residual block (IR Block) used in the ResNetArcFace architecture.

    Args:
        inplanes (int): Channel number of inputs.
        planes (int): Channel number of outputs.
        stride (int): Stride in convolution. Default: 1.
        downsample (nn.Module): The downsample module. Default: None.
        use_se (bool): Whether use the SEBlock (squeeze and excitation block). Default: True.
    """

    expansion = 1  # output channel expansion ratio

    def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True):
        super(IRBlock, self).__init__()
        self.bn0 = nn.BatchNorm2d(inplanes)
        self.conv1 = conv3x3(inplanes, inplanes)
        self.bn1 = nn.BatchNorm2d(inplanes)
        self.prelu = nn.PReLU()
        self.conv2 = conv3x3(inplanes, planes, stride)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride
        self.use_se = use_se
        if self.use_se:
            self.se = SEBlock(planes)

    def forward(self, x):
        residual = x
        out = self.bn0(x)
        out = self.conv1(out)
        out = self.bn1(out)
        out = self.prelu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        if self.use_se:
            out = self.se(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.prelu(out)

        return out


class Bottleneck(nn.Module):
    """Bottleneck block used in the ResNetArcFace architecture.

    Args:
        inplanes (int): Channel number of inputs.
        planes (int): Channel number of outputs.
        stride (int): Stride in convolution. Default: 1.
        downsample (nn.Module): The downsample module. Default: None.
    """

    expansion = 4  # output channel expansion ratio

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(
            planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(
            planes, planes * self.expansion, kernel_size=1, bias=False
        )
        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class SEBlock(nn.Module):
    """The squeeze-and-excitation block (SEBlock) used in the IRBlock.

    Args:
        channel (int): Channel number of inputs.
        reduction (int): Channel reduction ration. Default: 16.
    """

    def __init__(self, channel, reduction=16):
        super(SEBlock, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(
            1
        )  # pool to 1x1 without spatial information
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction),
            nn.PReLU(),
            nn.Linear(channel // reduction, channel),
            nn.Sigmoid(),
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y


class ResNetArcFace(nn.Module):
    """ArcFace with ResNet architectures.

    Ref: ArcFace: Additive Angular Margin Loss for Deep Face Recognition.

    Args:
        block (str): Block used in the ArcFace architecture.
        layers (tuple(int)): Block numbers in each layer.
        use_se (bool): Whether use the SEBlock (squeeze and excitation block). Default: True.
    """

    def __init__(self, block, layers, use_se=True):
        if block == "IRBlock":
            block = IRBlock
        self.inplanes = 64
        self.use_se = use_se
        super(ResNetArcFace, self).__init__()

        self.conv1 = nn.Conv2d(1, 64, kernel_size=3, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.prelu = nn.PReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.bn4 = nn.BatchNorm2d(512)
        self.dropout = nn.Dropout()
        self.fc5 = nn.Linear(512 * 8 * 8, 512)
        self.bn5 = nn.BatchNorm1d(512)

        # initialization
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.xavier_normal_(m.weight)
            elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                nn.init.constant_(m.bias, 0)

    def _make_layer(self, block, planes, num_blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(
                    self.inplanes,
                    planes * block.expansion,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                nn.BatchNorm2d(planes * block.expansion),
            )
        layers = []
        layers.append(
            block(self.inplanes, planes, stride, downsample, use_se=self.use_se)
        )
        self.inplanes = planes
        for _ in range(1, num_blocks):
            layers.append(block(self.inplanes, planes, use_se=self.use_se))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.prelu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.bn4(x)
        x = self.dropout(x)
        x = x.view(x.size(0), -1)
        x = self.fc5(x)
        x = self.bn5(x)

        return x

nodes/impl/pytorch/architecture/face/arcface_arch.py
vladmandic/sd-extension-chainner
"""
Modified from https://github.com/sczhou/CodeFormer
VQGAN code, adapted from the original created by the Unleashing Transformers authors:
https://github.com/samb-t/unleashing-transformers/blob/master/models/vqgan.py
This verison of the arch specifically was gathered from an old version of GFPGAN. If this is a problem, please contact me.
"""
import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from nodes.log import logger
from torch import Tensor


class VectorQuantizer(nn.Module):
    def __init__(self, codebook_size, emb_dim, beta):
        super(VectorQuantizer, self).__init__()
        self.codebook_size = codebook_size  # number of embeddings
        self.emb_dim = emb_dim  # dimension of embedding
        self.beta = beta  # commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2
        self.embedding = nn.Embedding(self.codebook_size, self.emb_dim)
        self.embedding.weight.data.uniform_(
            -1.0 / self.codebook_size, 1.0 / self.codebook_size
        )

    def forward(self, z):
        # reshape z -> (batch, height, width, channel) and flatten
        z = z.permute(0, 2, 3, 1).contiguous()
        z_flattened = z.view(-1, self.emb_dim)

        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
        d = (
            (z_flattened**2).sum(dim=1, keepdim=True)
            + (self.embedding.weight**2).sum(1)
            - 2 * torch.matmul(z_flattened, self.embedding.weight.t())
        )

        mean_distance = torch.mean(d)
        # find closest encodings
        # min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1)
        min_encoding_scores, min_encoding_indices = torch.topk(
            d, 1, dim=1, largest=False
        )
        # [0-1], higher score, higher confidence
        min_encoding_scores = torch.exp(-min_encoding_scores / 10)

        min_encodings = torch.zeros(
            min_encoding_indices.shape[0], self.codebook_size
        ).to(z)
        min_encodings.scatter_(1, min_encoding_indices, 1)

        # get quantized latent vectors
        z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape)
        # compute loss for embedding
        loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean(
            (z_q - z.detach()) ** 2
        )
        # preserve gradients
        z_q = z + (z_q - z).detach()

        # perplexity
        e_mean = torch.mean(min_encodings, dim=0)
        perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10)))
        # reshape back to match original input shape
        z_q = z_q.permute(0, 3, 1, 2).contiguous()

        return (
            z_q,
            loss,
            {
                "perplexity": perplexity,
                "min_encodings": min_encodings,
                "min_encoding_indices": min_encoding_indices,
                "min_encoding_scores": min_encoding_scores,
                "mean_distance": mean_distance,
            },
        )

    def get_codebook_feat(self, indices, shape):
        # input indices: batch*token_num -> (batch*token_num)*1
        # shape: batch, height, width, channel
        indices = indices.view(-1, 1)
        min_encodings = torch.zeros(indices.shape[0], self.codebook_size).to(indices)
        min_encodings.scatter_(1, indices, 1)
        # get quantized latent vectors
        z_q = torch.matmul(min_encodings.float(), self.embedding.weight)

        if shape is not None:  # reshape back to match original input shape
            z_q = z_q.view(shape).permute(0, 3, 1, 2).contiguous()

        return z_q


class GumbelQuantizer(nn.Module):
    def __init__(
        self,
        codebook_size,
        emb_dim,
        num_hiddens,
        straight_through=False,
        kl_weight=5e-4,
        temp_init=1.0,
    ):
        super().__init__()
        self.codebook_size = codebook_size  # number of embeddings
        self.emb_dim = emb_dim  # dimension of embedding
        self.straight_through = straight_through
        self.temperature = temp_init
        self.kl_weight = kl_weight
        self.proj = nn.Conv2d(
            num_hiddens, codebook_size, 1
        )  # projects last encoder layer to quantized logits
        self.embed = nn.Embedding(codebook_size, emb_dim)

    def forward(self, z):
        hard = self.straight_through if self.training else True

        logits = self.proj(z)

        soft_one_hot = F.gumbel_softmax(logits, tau=self.temperature, dim=1, hard=hard)

        z_q = torch.einsum("b n h w, n d -> b d h w", soft_one_hot, self.embed.weight)

        # + kl divergence to the prior loss
        qy = F.softmax(logits, dim=1)
        diff = (
            self.kl_weight
            * torch.sum(qy * torch.log(qy * self.codebook_size + 1e-10), dim=1).mean()
        )
        min_encoding_indices = soft_one_hot.argmax(dim=1)

        return z_q, diff, {"min_encoding_indices": min_encoding_indices}


class Downsample(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.conv = torch.nn.Conv2d(
            in_channels, in_channels, kernel_size=3, stride=2, padding=0
        )

    def forward(self, x):
        pad = (0, 1, 0, 1)
        x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
        x = self.conv(x)
        return x


class Upsample(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.conv = nn.Conv2d(
            in_channels, in_channels, kernel_size=3, stride=1, padding=1
        )

    def forward(self, x):
        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
        x = self.conv(x)

        return x


class AttnBlock(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.in_channels = in_channels

        self.norm = normalize(in_channels)
        self.q = torch.nn.Conv2d(
            in_channels, in_channels, kernel_size=1, stride=1, padding=0
        )
        self.k = torch.nn.Conv2d(
            in_channels, in_channels, kernel_size=1, stride=1, padding=0
        )
        self.v = torch.nn.Conv2d(
            in_channels, in_channels, kernel_size=1, stride=1, padding=0
        )
        self.proj_out = torch.nn.Conv2d(
            in_channels, in_channels, kernel_size=1, stride=1, padding=0
        )

    def forward(self, x):
        h_ = x
        h_ = self.norm(h_)
        q = self.q(h_)
        k = self.k(h_)
        v = self.v(h_)

        # compute attention
        b, c, h, w = q.shape
        q = q.reshape(b, c, h * w)
        q = q.permute(0, 2, 1)
        k = k.reshape(b, c, h * w)
        w_ = torch.bmm(q, k)
        w_ = w_ * (int(c) ** (-0.5))
        w_ = F.softmax(w_, dim=2)

        # attend to values
        v = v.reshape(b, c, h * w)
        w_ = w_.permute(0, 2, 1)
        h_ = torch.bmm(v, w_)
        h_ = h_.reshape(b, c, h, w)

        h_ = self.proj_out(h_)

        return x + h_


class Encoder(nn.Module):
    def __init__(
        self,
        in_channels,
        nf,
        out_channels,
        ch_mult,
        num_res_blocks,
        resolution,
        attn_resolutions,
    ):
        super().__init__()
        self.nf = nf
        self.num_resolutions = len(ch_mult)
        self.num_res_blocks = num_res_blocks
        self.resolution = resolution
        self.attn_resolutions = attn_resolutions

        curr_res = self.resolution
        in_ch_mult = (1,) + tuple(ch_mult)

        blocks = []
        # initial convultion
        blocks.append(nn.Conv2d(in_channels, nf, kernel_size=3, stride=1, padding=1))

        # residual and downsampling blocks, with attention on smaller res (16x16)
        for i in range(self.num_resolutions):
            block_in_ch = nf * in_ch_mult[i]
            block_out_ch = nf * ch_mult[i]
            for _ in range(self.num_res_blocks):
                blocks.append(ResBlock(block_in_ch, block_out_ch))
                block_in_ch = block_out_ch
                if curr_res in attn_resolutions:
                    blocks.append(AttnBlock(block_in_ch))

            if i != self.num_resolutions - 1:
                blocks.append(Downsample(block_in_ch))
                curr_res = curr_res // 2

        # non-local attention block
        blocks.append(ResBlock(block_in_ch, block_in_ch))  # type: ignore
        blocks.append(AttnBlock(block_in_ch))  # type: ignore
        blocks.append(ResBlock(block_in_ch, block_in_ch))  # type: ignore

        # normalise and convert to latent size
        blocks.append(normalize(block_in_ch))  # type: ignore
        blocks.append(
            nn.Conv2d(block_in_ch, out_channels, kernel_size=3, stride=1, padding=1)  # type: ignore
        )
        self.blocks = nn.ModuleList(blocks)

    def forward(self, x):
        for block in self.blocks:
            x = block(x)

        return x


class Generator(nn.Module):
    def __init__(self, nf, ch_mult, res_blocks, img_size, attn_resolutions, emb_dim):
        super().__init__()
        self.nf = nf
        self.ch_mult = ch_mult
        self.num_resolutions = len(self.ch_mult)
        self.num_res_blocks = res_blocks
        self.resolution = img_size
        self.attn_resolutions = attn_resolutions
        self.in_channels = emb_dim
        self.out_channels = 3
        block_in_ch = self.nf * self.ch_mult[-1]
        curr_res = self.resolution // 2 ** (self.num_resolutions - 1)

        blocks = []
        # initial conv
        blocks.append(
            nn.Conv2d(self.in_channels, block_in_ch, kernel_size=3, stride=1, padding=1)
        )

        # non-local attention block
        blocks.append(ResBlock(block_in_ch, block_in_ch))
        blocks.append(AttnBlock(block_in_ch))
        blocks.append(ResBlock(block_in_ch, block_in_ch))

        for i in reversed(range(self.num_resolutions)):
            block_out_ch = self.nf * self.ch_mult[i]

            for _ in range(self.num_res_blocks):
                blocks.append(ResBlock(block_in_ch, block_out_ch))
                block_in_ch = block_out_ch

                if curr_res in self.attn_resolutions:
                    blocks.append(AttnBlock(block_in_ch))

            if i != 0:
                blocks.append(Upsample(block_in_ch))
                curr_res = curr_res * 2

        blocks.append(normalize(block_in_ch))
        blocks.append(
            nn.Conv2d(
                block_in_ch, self.out_channels, kernel_size=3, stride=1, padding=1
            )
        )

        self.blocks = nn.ModuleList(blocks)

    def forward(self, x):
        for block in self.blocks:
            x = block(x)

        return x


class VQAutoEncoder(nn.Module):
    def __init__(
        self,
        img_size,
        nf,
        ch_mult,
        quantizer="nearest",
        res_blocks=2,
        attn_resolutions=None,
        codebook_size=1024,
        emb_dim=256,
        beta=0.25,
        gumbel_straight_through=False,
        gumbel_kl_weight=1e-8,
        model_path=None,
    ):
        if attn_resolutions is None:
            attn_resolutions = [16]
        super().__init__()
        self.in_channels = 3
        self.nf = nf
        self.n_blocks = res_blocks
        self.codebook_size = codebook_size
        self.embed_dim = emb_dim
        self.ch_mult = ch_mult
        self.resolution = img_size
        self.attn_resolutions = attn_resolutions
        self.quantizer_type = quantizer
        self.encoder = Encoder(
            self.in_channels,
            self.nf,
            self.embed_dim,
            self.ch_mult,
            self.n_blocks,
            self.resolution,
            self.attn_resolutions,
        )
        if self.quantizer_type == "nearest":
            self.beta = beta  # 0.25
            self.quantize = VectorQuantizer(
                self.codebook_size, self.embed_dim, self.beta
            )
        elif self.quantizer_type == "gumbel":
            self.gumbel_num_hiddens = emb_dim
            self.straight_through = gumbel_straight_through
            self.kl_weight = gumbel_kl_weight
            self.quantize = GumbelQuantizer(
                self.codebook_size,
                self.embed_dim,
                self.gumbel_num_hiddens,
                self.straight_through,
                self.kl_weight,
            )
        self.generator = Generator(
            nf, ch_mult, res_blocks, img_size, attn_resolutions, emb_dim
        )

        if model_path is not None:
            chkpt = torch.load(model_path, map_location="cpu")
            if "params_ema" in chkpt:
                self.load_state_dict(
                    torch.load(model_path, map_location="cpu")["params_ema"]
                )
                logger.info(f"chaiNNer: vqgan is loaded from: {model_path} [params_ema]")
            elif "params" in chkpt:
                self.load_state_dict(
                    torch.load(model_path, map_location="cpu")["params"]
                )
                logger.info(f"chaiNNer: vqgan is loaded from: {model_path} [params]")
            else:
                raise ValueError("Wrong params!")

    def forward(self, x):
        x = self.encoder(x)
        quant, codebook_loss, quant_stats = self.quantize(x)
        x = self.generator(quant)
        return x, codebook_loss, quant_stats


def calc_mean_std(feat, eps=1e-5):
    """Calculate mean and std for adaptive_instance_normalization.
    Args:
        feat (Tensor): 4D tensor.
        eps (float): A small value added to the variance to avoid
            divide-by-zero. Default: 1e-5.
    """
    size = feat.size()
    assert len(size) == 4, "The input feature should be 4D tensor."
    b, c = size[:2]
    feat_var = feat.view(b, c, -1).var(dim=2) + eps
    feat_std = feat_var.sqrt().view(b, c, 1, 1)
    feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1)
    return feat_mean, feat_std


def adaptive_instance_normalization(content_feat, style_feat):
    """Adaptive instance normalization.
    Adjust the reference features to have the similar color and illuminations
    as those in the degradate features.
    Args:
        content_feat (Tensor): The reference feature.
        style_feat (Tensor): The degradate features.
    """
    size = content_feat.size()
    style_mean, style_std = calc_mean_std(style_feat)
    content_mean, content_std = calc_mean_std(content_feat)
    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(
        size
    )
    return normalized_feat * style_std.expand(size) + style_mean.expand(size)


class PositionEmbeddingSine(nn.Module):
    """
    This is a more standard version of the position embedding, very similar to the one
    used by the Attention is all you need paper, generalized to work on images.
    """

    def __init__(
        self, num_pos_feats=64, temperature=10000, normalize=False, scale=None
    ):
        super().__init__()
        self.num_pos_feats = num_pos_feats
        self.temperature = temperature
        self.normalize = normalize
        if scale is not None and normalize is False:
            raise ValueError("normalize should be True if scale is passed")
        if scale is None:
            scale = 2 * math.pi
        self.scale = scale

    def forward(self, x, mask=None):
        if mask is None:
            mask = torch.zeros(
                (x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool
            )
        not_mask = ~mask  # pylint: disable=invalid-unary-operand-type
        y_embed = not_mask.cumsum(1, dtype=torch.float32)
        x_embed = not_mask.cumsum(2, dtype=torch.float32)
        if self.normalize:
            eps = 1e-6
            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale

        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)

        pos_x = x_embed[:, :, :, None] / dim_t
        pos_y = y_embed[:, :, :, None] / dim_t
        pos_x = torch.stack(
            (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4
        ).flatten(3)
        pos_y = torch.stack(
            (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4
        ).flatten(3)
        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
        return pos


def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(f"activation should be relu/gelu, not {activation}.")


class TransformerSALayer(nn.Module):
    def __init__(
        self, embed_dim, nhead=8, dim_mlp=2048, dropout=0.0, activation="gelu"
    ):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(embed_dim, nhead, dropout=dropout)
        # Implementation of Feedforward model - MLP
        self.linear1 = nn.Linear(embed_dim, dim_mlp)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_mlp, embed_dim)

        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward(
        self,
        tgt,
        tgt_mask: Optional[Tensor] = None,
        tgt_key_padding_mask: Optional[Tensor] = None,
        query_pos: Optional[Tensor] = None,
    ):
        # self attention
        tgt2 = self.norm1(tgt)
        q = k = self.with_pos_embed(tgt2, query_pos)
        tgt2 = self.self_attn(
            q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
        )[0]
        tgt = tgt + self.dropout1(tgt2)

        # ffn
        tgt2 = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout2(tgt2)
        return tgt


def normalize(in_channels):
    return torch.nn.GroupNorm(
        num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
    )


@torch.jit.script  # type: ignore
def swish(x):
    return x * torch.sigmoid(x)


class ResBlock(nn.Module):
    def __init__(self, in_channels, out_channels=None):
        super(ResBlock, self).__init__()
        self.in_channels = in_channels
        self.out_channels = in_channels if out_channels is None else out_channels
        self.norm1 = normalize(in_channels)
        self.conv1 = nn.Conv2d(
            in_channels, out_channels, kernel_size=3, stride=1, padding=1  # type: ignore
        )
        self.norm2 = normalize(out_channels)
        self.conv2 = nn.Conv2d(
            out_channels, out_channels, kernel_size=3, stride=1, padding=1  # type: ignore
        )
        if self.in_channels != self.out_channels:
            self.conv_out = nn.Conv2d(
                in_channels, out_channels, kernel_size=1, stride=1, padding=0  # type: ignore
            )

    def forward(self, x_in):
        x = x_in
        x = self.norm1(x)
        x = swish(x)
        x = self.conv1(x)
        x = self.norm2(x)
        x = swish(x)
        x = self.conv2(x)
        if self.in_channels != self.out_channels:
            x_in = self.conv_out(x_in)

        return x + x_in


class Fuse_sft_block(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.encode_enc = ResBlock(2 * in_ch, out_ch)

        self.scale = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, True),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
        )

        self.shift = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, True),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
        )

    def forward(self, enc_feat, dec_feat, w=1):
        enc_feat = self.encode_enc(torch.cat([enc_feat, dec_feat], dim=1))
        scale = self.scale(enc_feat)
        shift = self.shift(enc_feat)
        residual = w * (dec_feat * scale + shift)
        out = dec_feat + residual
        return out


class CodeFormer(VQAutoEncoder):
    def __init__(self, state_dict):
        dim_embd = 512
        n_head = 8
        n_layers = 9
        codebook_size = 1024
        latent_size = 256
        connect_list = ["32", "64", "128", "256"]
        fix_modules = ["quantize", "generator"]

        # This is just a guess as I only have one model to look at
        position_emb = state_dict["position_emb"]
        dim_embd = position_emb.shape[1]
        latent_size = position_emb.shape[0]

        try:
            n_layers = len(
                {x.split(".")[1] for x in state_dict.keys() if "ft_layers" in x}
            )
        except:
            pass

        codebook_size = state_dict["quantize.embedding.weight"].shape[0]

        # This is also just another guess
        n_head_exp = (
            state_dict["ft_layers.0.self_attn.in_proj_weight"].shape[0] // dim_embd
        )
        n_head = 2**n_head_exp

        in_nc = state_dict["encoder.blocks.0.weight"].shape[1]

        self.model_arch = "CodeFormer"
        self.sub_type = "Face SR"
        self.scale = 8
        self.in_nc = in_nc
        self.out_nc = in_nc

        self.state = state_dict

        self.supports_fp16 = False
        self.supports_bf16 = True
        self.min_size_restriction = 16

        super(CodeFormer, self).__init__(
            512, 64, [1, 2, 2, 4, 4, 8], "nearest", 2, [16], codebook_size
        )

        if fix_modules is not None:
            for module in fix_modules:
                for param in getattr(self, module).parameters():
                    param.requires_grad = False

        self.connect_list = connect_list
        self.n_layers = n_layers
        self.dim_embd = dim_embd
        self.dim_mlp = dim_embd * 2

        self.position_emb = nn.Parameter(torch.zeros(latent_size, self.dim_embd))  # type: ignore
        self.feat_emb = nn.Linear(256, self.dim_embd)

        # transformer
        self.ft_layers = nn.Sequential(
            *[
                TransformerSALayer(
                    embed_dim=dim_embd, nhead=n_head, dim_mlp=self.dim_mlp, dropout=0.0
                )
                for _ in range(self.n_layers)
            ]
        )

        # logits_predict head
        self.idx_pred_layer = nn.Sequential(
            nn.LayerNorm(dim_embd), nn.Linear(dim_embd, codebook_size, bias=False)
        )

        self.channels = {
            "16": 512,
            "32": 256,
            "64": 256,
            "128": 128,
            "256": 128,
            "512": 64,
        }

        # after second residual block for > 16, before attn layer for ==16
        self.fuse_encoder_block = {
            "512": 2,
            "256": 5,
            "128": 8,
            "64": 11,
            "32": 14,
            "16": 18,
        }
        # after first residual block for > 16, before attn layer for ==16
        self.fuse_generator_block = {
            "16": 6,
            "32": 9,
            "64": 12,
            "128": 15,
            "256": 18,
            "512": 21,
        }

        # fuse_convs_dict
        self.fuse_convs_dict = nn.ModuleDict()
        for f_size in self.connect_list:
            in_ch = self.channels[f_size]
            self.fuse_convs_dict[f_size] = Fuse_sft_block(in_ch, in_ch)

        self.load_state_dict(state_dict)

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def forward(self, x, weight=0.5, **kwargs):
        detach_16 = True
        code_only = False
        adain = True
        # ################### Encoder #####################
        enc_feat_dict = {}
        out_list = [self.fuse_encoder_block[f_size] for f_size in self.connect_list]
        for i, block in enumerate(self.encoder.blocks):
            x = block(x)
            if i in out_list:
                enc_feat_dict[str(x.shape[-1])] = x.clone()

        lq_feat = x
        # ################# Transformer ###################
        # quant_feat, codebook_loss, quant_stats = self.quantize(lq_feat)
        pos_emb = self.position_emb.unsqueeze(1).repeat(1, x.shape[0], 1)
        # BCHW -> BC(HW) -> (HW)BC
        feat_emb = self.feat_emb(lq_feat.flatten(2).permute(2, 0, 1))
        query_emb = feat_emb
        # Transformer encoder
        for layer in self.ft_layers:
            query_emb = layer(query_emb, query_pos=pos_emb)

        # output logits
        logits = self.idx_pred_layer(query_emb)  # (hw)bn
        logits = logits.permute(1, 0, 2)  # (hw)bn -> b(hw)n

        if code_only:  # for training stage II
            # logits doesn't need softmax before cross_entropy loss
            return logits, lq_feat

        # ################# Quantization ###################
        # if self.training:
        #     quant_feat = torch.einsum('btn,nc->btc', [soft_one_hot, self.quantize.embedding.weight])
        #     # b(hw)c -> bc(hw) -> bchw
        #     quant_feat = quant_feat.permute(0,2,1).view(lq_feat.shape)
        # ------------
        soft_one_hot = F.softmax(logits, dim=2)
        _, top_idx = torch.topk(soft_one_hot, 1, dim=2)
        quant_feat = self.quantize.get_codebook_feat(
            top_idx, shape=[x.shape[0], 16, 16, 256]  # type: ignore
        )
        # preserve gradients
        # quant_feat = lq_feat + (quant_feat - lq_feat).detach()

        if detach_16:
            quant_feat = quant_feat.detach()  # for training stage III
        if adain:
            quant_feat = adaptive_instance_normalization(quant_feat, lq_feat)

        # ################## Generator ####################
        x = quant_feat
        fuse_list = [self.fuse_generator_block[f_size] for f_size in self.connect_list]

        for i, block in enumerate(self.generator.blocks):
            x = block(x)
            if i in fuse_list:  # fuse after i-th block
                f_size = str(x.shape[-1])
                if weight > 0:
                    x = self.fuse_convs_dict[f_size](
                        enc_feat_dict[f_size].detach(), x, weight
                    )
        out = x
        # logits doesn't need softmax before cross_entropy loss
        # return out, logits, lq_feat
        return out, logits

nodes/impl/pytorch/architecture/face/codeformer.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# type: ignore
# modify from https://github.com/rosinality/stylegan2-pytorch/blob/master/op/fused_act.py

import torch
from torch import nn
from torch.autograd import Function

fused_act_ext = None


class FusedLeakyReLUFunctionBackward(Function):
    @staticmethod
    def forward(ctx, grad_output, out, negative_slope, scale):
        ctx.save_for_backward(out)
        ctx.negative_slope = negative_slope
        ctx.scale = scale

        empty = grad_output.new_empty(0)

        grad_input = fused_act_ext.fused_bias_act(
            grad_output, empty, out, 3, 1, negative_slope, scale
        )

        dim = [0]

        if grad_input.ndim > 2:
            dim += list(range(2, grad_input.ndim))

        grad_bias = grad_input.sum(dim).detach()

        return grad_input, grad_bias

    @staticmethod
    def backward(ctx, gradgrad_input, gradgrad_bias):
        (out,) = ctx.saved_tensors
        gradgrad_out = fused_act_ext.fused_bias_act(
            gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale
        )

        return gradgrad_out, None, None, None


class FusedLeakyReLUFunction(Function):
    @staticmethod
    def forward(ctx, input, bias, negative_slope, scale):
        empty = input.new_empty(0)
        out = fused_act_ext.fused_bias_act(
            input, bias, empty, 3, 0, negative_slope, scale
        )
        ctx.save_for_backward(out)
        ctx.negative_slope = negative_slope
        ctx.scale = scale

        return out

    @staticmethod
    def backward(ctx, grad_output):
        (out,) = ctx.saved_tensors

        grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply(
            grad_output, out, ctx.negative_slope, ctx.scale
        )

        return grad_input, grad_bias, None, None


class FusedLeakyReLU(nn.Module):
    def __init__(self, channel, negative_slope=0.2, scale=2**0.5):
        super().__init__()

        self.bias = nn.Parameter(torch.zeros(channel))
        self.negative_slope = negative_slope
        self.scale = scale

    def forward(self, input):
        return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)


def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2**0.5):
    return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale)

nodes/impl/pytorch/architecture/face/fused_act.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# type: ignore
import math
import random

import torch
from torch import nn

from .gfpganv1_arch import ResUpBlock
from .stylegan2_bilinear_arch import (
    ConvLayer,
    EqualConv2d,
    EqualLinear,
    ResBlock,
    ScaledLeakyReLU,
    StyleGAN2GeneratorBilinear,
)


class StyleGAN2GeneratorBilinearSFT(StyleGAN2GeneratorBilinear):
    """StyleGAN2 Generator with SFT modulation (Spatial Feature Transform).
    It is the bilinear version. It does not use the complicated UpFirDnSmooth function that is not friendly for
    deployment. It can be easily converted to the clean version: StyleGAN2GeneratorCSFT.
    Args:
        out_size (int): The spatial size of outputs.
        num_style_feat (int): Channel number of style features. Default: 512.
        num_mlp (int): Layer number of MLP style layers. Default: 8.
        channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
        lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01.
        narrow (float): The narrow ratio for channels. Default: 1.
        sft_half (bool): Whether to apply SFT on half of the input channels. Default: False.
    """

    def __init__(
        self,
        out_size,
        num_style_feat=512,
        num_mlp=8,
        channel_multiplier=2,
        lr_mlp=0.01,
        narrow=1,
        sft_half=False,
    ):
        super(StyleGAN2GeneratorBilinearSFT, self).__init__(
            out_size,
            num_style_feat=num_style_feat,
            num_mlp=num_mlp,
            channel_multiplier=channel_multiplier,
            lr_mlp=lr_mlp,
            narrow=narrow,
        )
        self.sft_half = sft_half

    def forward(
        self,
        styles,
        conditions,
        input_is_latent=False,
        noise=None,
        randomize_noise=True,
        truncation=1,
        truncation_latent=None,
        inject_index=None,
        return_latents=False,
    ):
        """Forward function for StyleGAN2GeneratorBilinearSFT.
        Args:
            styles (list[Tensor]): Sample codes of styles.
            conditions (list[Tensor]): SFT conditions to generators.
            input_is_latent (bool): Whether input is latent style. Default: False.
            noise (Tensor | None): Input noise or None. Default: None.
            randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True.
            truncation (float): The truncation ratio. Default: 1.
            truncation_latent (Tensor | None): The truncation latent tensor. Default: None.
            inject_index (int | None): The injection index for mixing noise. Default: None.
            return_latents (bool): Whether to return style latents. Default: False.
        """
        # style codes -> latents with Style MLP layer
        if not input_is_latent:
            styles = [self.style_mlp(s) for s in styles]
        # noises
        if noise is None:
            if randomize_noise:
                noise = [None] * self.num_layers  # for each style conv layer
            else:  # use the stored noise
                noise = [
                    getattr(self.noises, f"noise{i}") for i in range(self.num_layers)
                ]
        # style truncation
        if truncation < 1:
            style_truncation = []
            for style in styles:
                style_truncation.append(
                    truncation_latent + truncation * (style - truncation_latent)
                )
            styles = style_truncation
        # get style latents with injection
        if len(styles) == 1:
            inject_index = self.num_latent

            if styles[0].ndim < 3:
                # repeat latent code for all the layers
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            else:  # used for encoder with different latent code for each layer
                latent = styles[0]
        elif len(styles) == 2:  # mixing noises
            if inject_index is None:
                inject_index = random.randint(1, self.num_latent - 1)
            latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            latent2 = (
                styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1)
            )
            latent = torch.cat([latent1, latent2], 1)

        # main generation
        out = self.constant_input(latent.shape[0])
        out = self.style_conv1(out, latent[:, 0], noise=noise[0])
        skip = self.to_rgb1(out, latent[:, 1])

        i = 1
        for conv1, conv2, noise1, noise2, to_rgb in zip(
            self.style_convs[::2],
            self.style_convs[1::2],
            noise[1::2],
            noise[2::2],
            self.to_rgbs,
        ):
            out = conv1(out, latent[:, i], noise=noise1)

            # the conditions may have fewer levels
            if i < len(conditions):
                # SFT part to combine the conditions
                if self.sft_half:  # only apply SFT to half of the channels
                    out_same, out_sft = torch.split(out, int(out.size(1) // 2), dim=1)
                    out_sft = out_sft * conditions[i - 1] + conditions[i]
                    out = torch.cat([out_same, out_sft], dim=1)
                else:  # apply SFT to all the channels
                    out = out * conditions[i - 1] + conditions[i]

            out = conv2(out, latent[:, i + 1], noise=noise2)
            skip = to_rgb(out, latent[:, i + 2], skip)  # feature back to the rgb space
            i += 2

        image = skip

        if return_latents:
            return image, latent
        else:
            return image, None


class GFPGANBilinear(nn.Module):
    """The GFPGAN architecture: Unet + StyleGAN2 decoder with SFT.
    It is the bilinear version and it does not use the complicated UpFirDnSmooth function that is not friendly for
    deployment. It can be easily converted to the clean version: GFPGANv1Clean.
    Ref: GFP-GAN: Towards Real-World Blind Face Restoration with Generative Facial Prior.
    Args:
        out_size (int): The spatial size of outputs.
        num_style_feat (int): Channel number of style features. Default: 512.
        channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
        decoder_load_path (str): The path to the pre-trained decoder model (usually, the StyleGAN2). Default: None.
        fix_decoder (bool): Whether to fix the decoder. Default: True.
        num_mlp (int): Layer number of MLP style layers. Default: 8.
        lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01.
        input_is_latent (bool): Whether input is latent style. Default: False.
        different_w (bool): Whether to use different latent w for different layers. Default: False.
        narrow (float): The narrow ratio for channels. Default: 1.
        sft_half (bool): Whether to apply SFT on half of the input channels. Default: False.
    """

    def __init__(
        self,
        out_size,
        num_style_feat=512,
        channel_multiplier=1,
        decoder_load_path=None,
        fix_decoder=True,
        # for stylegan decoder
        num_mlp=8,
        lr_mlp=0.01,
        input_is_latent=False,
        different_w=False,
        narrow=1,
        sft_half=False,
    ):
        super(GFPGANBilinear, self).__init__()
        self.input_is_latent = input_is_latent
        self.different_w = different_w
        self.num_style_feat = num_style_feat
        self.min_size_restriction = 512

        unet_narrow = narrow * 0.5  # by default, use a half of input channels
        channels = {
            "4": int(512 * unet_narrow),
            "8": int(512 * unet_narrow),
            "16": int(512 * unet_narrow),
            "32": int(512 * unet_narrow),
            "64": int(256 * channel_multiplier * unet_narrow),
            "128": int(128 * channel_multiplier * unet_narrow),
            "256": int(64 * channel_multiplier * unet_narrow),
            "512": int(32 * channel_multiplier * unet_narrow),
            "1024": int(16 * channel_multiplier * unet_narrow),
        }

        self.log_size = int(math.log(out_size, 2))
        first_out_size = 2 ** (int(math.log(out_size, 2)))

        self.conv_body_first = ConvLayer(
            3, channels[f"{first_out_size}"], 1, bias=True, activate=True
        )

        # downsample
        in_channels = channels[f"{first_out_size}"]
        self.conv_body_down = nn.ModuleList()
        for i in range(self.log_size, 2, -1):
            out_channels = channels[f"{2**(i - 1)}"]
            self.conv_body_down.append(ResBlock(in_channels, out_channels))
            in_channels = out_channels

        self.final_conv = ConvLayer(
            in_channels, channels["4"], 3, bias=True, activate=True
        )

        # upsample
        in_channels = channels["4"]
        self.conv_body_up = nn.ModuleList()
        for i in range(3, self.log_size + 1):
            out_channels = channels[f"{2**i}"]
            self.conv_body_up.append(ResUpBlock(in_channels, out_channels))
            in_channels = out_channels

        # to RGB
        self.toRGB = nn.ModuleList()
        for i in range(3, self.log_size + 1):
            self.toRGB.append(
                EqualConv2d(
                    channels[f"{2**i}"],
                    3,
                    1,
                    stride=1,
                    padding=0,
                    bias=True,
                    bias_init_val=0,
                )
            )

        if different_w:
            linear_out_channel = (int(math.log(out_size, 2)) * 2 - 2) * num_style_feat
        else:
            linear_out_channel = num_style_feat

        self.final_linear = EqualLinear(
            channels["4"] * 4 * 4,
            linear_out_channel,
            bias=True,
            bias_init_val=0,
            lr_mul=1,
            activation=None,
        )

        # the decoder: stylegan2 generator with SFT modulations
        self.stylegan_decoder = StyleGAN2GeneratorBilinearSFT(
            out_size=out_size,
            num_style_feat=num_style_feat,
            num_mlp=num_mlp,
            channel_multiplier=channel_multiplier,
            lr_mlp=lr_mlp,
            narrow=narrow,
            sft_half=sft_half,
        )

        # load pre-trained stylegan2 model if necessary
        if decoder_load_path:
            self.stylegan_decoder.load_state_dict(
                torch.load(
                    decoder_load_path, map_location=lambda storage, loc: storage
                )["params_ema"]
            )
        # fix decoder without updating params
        if fix_decoder:
            for _, param in self.stylegan_decoder.named_parameters():
                param.requires_grad = False

        # for SFT modulations (scale and shift)
        self.condition_scale = nn.ModuleList()
        self.condition_shift = nn.ModuleList()
        for i in range(3, self.log_size + 1):
            out_channels = channels[f"{2**i}"]
            if sft_half:
                sft_out_channels = out_channels
            else:
                sft_out_channels = out_channels * 2
            self.condition_scale.append(
                nn.Sequential(
                    EqualConv2d(
                        out_channels,
                        out_channels,
                        3,
                        stride=1,
                        padding=1,
                        bias=True,
                        bias_init_val=0,
                    ),
                    ScaledLeakyReLU(0.2),
                    EqualConv2d(
                        out_channels,
                        sft_out_channels,
                        3,
                        stride=1,
                        padding=1,
                        bias=True,
                        bias_init_val=1,
                    ),
                )
            )
            self.condition_shift.append(
                nn.Sequential(
                    EqualConv2d(
                        out_channels,
                        out_channels,
                        3,
                        stride=1,
                        padding=1,
                        bias=True,
                        bias_init_val=0,
                    ),
                    ScaledLeakyReLU(0.2),
                    EqualConv2d(
                        out_channels,
                        sft_out_channels,
                        3,
                        stride=1,
                        padding=1,
                        bias=True,
                        bias_init_val=0,
                    ),
                )
            )

    def forward(self, x, return_latents=False, return_rgb=True, randomize_noise=True):
        """Forward function for GFPGANBilinear.
        Args:
            x (Tensor): Input images.
            return_latents (bool): Whether to return style latents. Default: False.
            return_rgb (bool): Whether return intermediate rgb images. Default: True.
            randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True.
        """
        conditions = []
        unet_skips = []
        out_rgbs = []

        # encoder
        feat = self.conv_body_first(x)
        for i in range(self.log_size - 2):
            feat = self.conv_body_down[i](feat)
            unet_skips.insert(0, feat)

        feat = self.final_conv(feat)

        # style code
        style_code = self.final_linear(feat.view(feat.size(0), -1))
        if self.different_w:
            style_code = style_code.view(style_code.size(0), -1, self.num_style_feat)

        # decode
        for i in range(self.log_size - 2):
            # add unet skip
            feat = feat + unet_skips[i]
            # ResUpLayer
            feat = self.conv_body_up[i](feat)
            # generate scale and shift for SFT layers
            scale = self.condition_scale[i](feat)
            conditions.append(scale.clone())
            shift = self.condition_shift[i](feat)
            conditions.append(shift.clone())
            # generate rgb images
            if return_rgb:
                out_rgbs.append(self.toRGB[i](feat))

        # decoder
        image, _ = self.stylegan_decoder(
            [style_code],
            conditions,
            return_latents=return_latents,
            input_is_latent=self.input_is_latent,
            randomize_noise=randomize_noise,
        )

        return image, out_rgbs

nodes/impl/pytorch/architecture/face/gfpgan_bilinear_arch.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# type: ignore
import math
import random

import torch
from torch import nn
from torch.nn import functional as F

from .fused_act import FusedLeakyReLU
from .stylegan2_arch import (
    ConvLayer,
    EqualConv2d,
    EqualLinear,
    ResBlock,
    ScaledLeakyReLU,
    StyleGAN2Generator,
)


class StyleGAN2GeneratorSFT(StyleGAN2Generator):
    """StyleGAN2 Generator with SFT modulation (Spatial Feature Transform).
    Args:
        out_size (int): The spatial size of outputs.
        num_style_feat (int): Channel number of style features. Default: 512.
        num_mlp (int): Layer number of MLP style layers. Default: 8.
        channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
        resample_kernel (list[int]): A list indicating the 1D resample kernel magnitude. A cross production will be
            applied to extent 1D resample kernel to 2D resample kernel. Default: (1, 3, 3, 1).
        lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01.
        narrow (float): The narrow ratio for channels. Default: 1.
        sft_half (bool): Whether to apply SFT on half of the input channels. Default: False.
    """

    def __init__(
        self,
        out_size,
        num_style_feat=512,
        num_mlp=8,
        channel_multiplier=2,
        resample_kernel=(1, 3, 3, 1),
        lr_mlp=0.01,
        narrow=1,
        sft_half=False,
    ):
        super(StyleGAN2GeneratorSFT, self).__init__(
            out_size,
            num_style_feat=num_style_feat,
            num_mlp=num_mlp,
            channel_multiplier=channel_multiplier,
            resample_kernel=resample_kernel,
            lr_mlp=lr_mlp,
            narrow=narrow,
        )
        self.sft_half = sft_half

    def forward(
        self,
        styles,
        conditions,
        input_is_latent=False,
        noise=None,
        randomize_noise=True,
        truncation=1,
        truncation_latent=None,
        inject_index=None,
        return_latents=False,
    ):
        """Forward function for StyleGAN2GeneratorSFT.
        Args:
            styles (list[Tensor]): Sample codes of styles.
            conditions (list[Tensor]): SFT conditions to generators.
            input_is_latent (bool): Whether input is latent style. Default: False.
            noise (Tensor | None): Input noise or None. Default: None.
            randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True.
            truncation (float): The truncation ratio. Default: 1.
            truncation_latent (Tensor | None): The truncation latent tensor. Default: None.
            inject_index (int | None): The injection index for mixing noise. Default: None.
            return_latents (bool): Whether to return style latents. Default: False.
        """
        # style codes -> latents with Style MLP layer
        if not input_is_latent:
            styles = [self.style_mlp(s) for s in styles]
        # noises
        if noise is None:
            if randomize_noise:
                noise = [None] * self.num_layers  # for each style conv layer
            else:  # use the stored noise
                noise = [
                    getattr(self.noises, f"noise{i}") for i in range(self.num_layers)
                ]
        # style truncation
        if truncation < 1:
            style_truncation = []
            for style in styles:
                style_truncation.append(
                    truncation_latent + truncation * (style - truncation_latent)
                )
            styles = style_truncation
        # get style latents with injection
        if len(styles) == 1:
            inject_index = self.num_latent

            if styles[0].ndim < 3:
                # repeat latent code for all the layers
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            else:  # used for encoder with different latent code for each layer
                latent = styles[0]
        elif len(styles) == 2:  # mixing noises
            if inject_index is None:
                inject_index = random.randint(1, self.num_latent - 1)
            latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            latent2 = (
                styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1)
            )
            latent = torch.cat([latent1, latent2], 1)

        # main generation
        out = self.constant_input(latent.shape[0])
        out = self.style_conv1(out, latent[:, 0], noise=noise[0])
        skip = self.to_rgb1(out, latent[:, 1])

        i = 1
        for conv1, conv2, noise1, noise2, to_rgb in zip(
            self.style_convs[::2],
            self.style_convs[1::2],
            noise[1::2],
            noise[2::2],
            self.to_rgbs,
        ):
            out = conv1(out, latent[:, i], noise=noise1)

            # the conditions may have fewer levels
            if i < len(conditions):
                # SFT part to combine the conditions
                if self.sft_half:  # only apply SFT to half of the channels
                    out_same, out_sft = torch.split(out, int(out.size(1) // 2), dim=1)
                    out_sft = out_sft * conditions[i - 1] + conditions[i]
                    out = torch.cat([out_same, out_sft], dim=1)
                else:  # apply SFT to all the channels
                    out = out * conditions[i - 1] + conditions[i]

            out = conv2(out, latent[:, i + 1], noise=noise2)
            skip = to_rgb(out, latent[:, i + 2], skip)  # feature back to the rgb space
            i += 2

        image = skip

        if return_latents:
            return image, latent
        else:
            return image, None


class ConvUpLayer(nn.Module):
    """Convolutional upsampling layer. It uses bilinear upsampler + Conv.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Size of the convolving kernel.
        stride (int): Stride of the convolution. Default: 1
        padding (int): Zero-padding added to both sides of the input. Default: 0.
        bias (bool): If ``True``, adds a learnable bias to the output. Default: ``True``.
        bias_init_val (float): Bias initialized value. Default: 0.
        activate (bool): Whether use activateion. Default: True.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        padding=0,
        bias=True,
        bias_init_val=0,
        activate=True,
    ):
        super(ConvUpLayer, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        # self.scale is used to scale the convolution weights, which is related to the common initializations.
        self.scale = 1 / math.sqrt(in_channels * kernel_size**2)

        self.weight = nn.Parameter(
            torch.randn(out_channels, in_channels, kernel_size, kernel_size)
        )

        if bias and not activate:
            self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val))
        else:
            self.register_parameter("bias", None)

        # activation
        if activate:
            if bias:
                self.activation = FusedLeakyReLU(out_channels)
            else:
                self.activation = ScaledLeakyReLU(0.2)
        else:
            self.activation = None

    def forward(self, x):
        # bilinear upsample
        out = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False)
        # conv
        out = F.conv2d(
            out,
            self.weight * self.scale,
            bias=self.bias,
            stride=self.stride,
            padding=self.padding,
        )
        # activation
        if self.activation is not None:
            out = self.activation(out)
        return out


class ResUpBlock(nn.Module):
    """Residual block with upsampling.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
    """

    def __init__(self, in_channels, out_channels):
        super(ResUpBlock, self).__init__()

        self.conv1 = ConvLayer(in_channels, in_channels, 3, bias=True, activate=True)
        self.conv2 = ConvUpLayer(
            in_channels, out_channels, 3, stride=1, padding=1, bias=True, activate=True
        )
        self.skip = ConvUpLayer(
            in_channels, out_channels, 1, bias=False, activate=False
        )

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        skip = self.skip(x)
        out = (out + skip) / math.sqrt(2)
        return out


class GFPGANv1(nn.Module):
    """The GFPGAN architecture: Unet + StyleGAN2 decoder with SFT.
    Ref: GFP-GAN: Towards Real-World Blind Face Restoration with Generative Facial Prior.
    Args:
        out_size (int): The spatial size of outputs.
        num_style_feat (int): Channel number of style features. Default: 512.
        channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
        resample_kernel (list[int]): A list indicating the 1D resample kernel magnitude. A cross production will be
            applied to extent 1D resample kernel to 2D resample kernel. Default: (1, 3, 3, 1).
        decoder_load_path (str): The path to the pre-trained decoder model (usually, the StyleGAN2). Default: None.
        fix_decoder (bool): Whether to fix the decoder. Default: True.
        num_mlp (int): Layer number of MLP style layers. Default: 8.
        lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01.
        input_is_latent (bool): Whether input is latent style. Default: False.
        different_w (bool): Whether to use different latent w for different layers. Default: False.
        narrow (float): The narrow ratio for channels. Default: 1.
        sft_half (bool): Whether to apply SFT on half of the input channels. Default: False.
    """

    def __init__(
        self,
        out_size,
        num_style_feat=512,
        channel_multiplier=1,
        resample_kernel=(1, 3, 3, 1),
        decoder_load_path=None,
        fix_decoder=True,
        # for stylegan decoder
        num_mlp=8,
        lr_mlp=0.01,
        input_is_latent=False,
        different_w=False,
        narrow=1,
        sft_half=False,
    ):
        super(GFPGANv1, self).__init__()
        self.input_is_latent = input_is_latent
        self.different_w = different_w
        self.num_style_feat = num_style_feat

        unet_narrow = narrow * 0.5  # by default, use a half of input channels
        channels = {
            "4": int(512 * unet_narrow),
            "8": int(512 * unet_narrow),
            "16": int(512 * unet_narrow),
            "32": int(512 * unet_narrow),
            "64": int(256 * channel_multiplier * unet_narrow),
            "128": int(128 * channel_multiplier * unet_narrow),
            "256": int(64 * channel_multiplier * unet_narrow),
            "512": int(32 * channel_multiplier * unet_narrow),
            "1024": int(16 * channel_multiplier * unet_narrow),
        }

        self.log_size = int(math.log(out_size, 2))
        first_out_size = 2 ** (int(math.log(out_size, 2)))

        self.conv_body_first = ConvLayer(
            3, channels[f"{first_out_size}"], 1, bias=True, activate=True
        )

        # downsample
        in_channels = channels[f"{first_out_size}"]
        self.conv_body_down = nn.ModuleList()
        for i in range(self.log_size, 2, -1):
            out_channels = channels[f"{2**(i - 1)}"]
            self.conv_body_down.append(
                ResBlock(in_channels, out_channels, resample_kernel)
            )
            in_channels = out_channels

        self.final_conv = ConvLayer(
            in_channels, channels["4"], 3, bias=True, activate=True
        )

        # upsample
        in_channels = channels["4"]
        self.conv_body_up = nn.ModuleList()
        for i in range(3, self.log_size + 1):
            out_channels = channels[f"{2**i}"]
            self.conv_body_up.append(ResUpBlock(in_channels, out_channels))
            in_channels = out_channels

        # to RGB
        self.toRGB = nn.ModuleList()
        for i in range(3, self.log_size + 1):
            self.toRGB.append(
                EqualConv2d(
                    channels[f"{2**i}"],
                    3,
                    1,
                    stride=1,
                    padding=0,
                    bias=True,
                    bias_init_val=0,
                )
            )

        if different_w:
            linear_out_channel = (int(math.log(out_size, 2)) * 2 - 2) * num_style_feat
        else:
            linear_out_channel = num_style_feat

        self.final_linear = EqualLinear(
            channels["4"] * 4 * 4,
            linear_out_channel,
            bias=True,
            bias_init_val=0,
            lr_mul=1,
            activation=None,
        )

        # the decoder: stylegan2 generator with SFT modulations
        self.stylegan_decoder = StyleGAN2GeneratorSFT(
            out_size=out_size,
            num_style_feat=num_style_feat,
            num_mlp=num_mlp,
            channel_multiplier=channel_multiplier,
            resample_kernel=resample_kernel,
            lr_mlp=lr_mlp,
            narrow=narrow,
            sft_half=sft_half,
        )

        # load pre-trained stylegan2 model if necessary
        if decoder_load_path:
            self.stylegan_decoder.load_state_dict(
                torch.load(
                    decoder_load_path, map_location=lambda storage, loc: storage
                )["params_ema"]
            )
        # fix decoder without updating params
        if fix_decoder:
            for _, param in self.stylegan_decoder.named_parameters():
                param.requires_grad = False

        # for SFT modulations (scale and shift)
        self.condition_scale = nn.ModuleList()
        self.condition_shift = nn.ModuleList()
        for i in range(3, self.log_size + 1):
            out_channels = channels[f"{2**i}"]
            if sft_half:
                sft_out_channels = out_channels
            else:
                sft_out_channels = out_channels * 2
            self.condition_scale.append(
                nn.Sequential(
                    EqualConv2d(
                        out_channels,
                        out_channels,
                        3,
                        stride=1,
                        padding=1,
                        bias=True,
                        bias_init_val=0,
                    ),
                    ScaledLeakyReLU(0.2),
                    EqualConv2d(
                        out_channels,
                        sft_out_channels,
                        3,
                        stride=1,
                        padding=1,
                        bias=True,
                        bias_init_val=1,
                    ),
                )
            )
            self.condition_shift.append(
                nn.Sequential(
                    EqualConv2d(
                        out_channels,
                        out_channels,
                        3,
                        stride=1,
                        padding=1,
                        bias=True,
                        bias_init_val=0,
                    ),
                    ScaledLeakyReLU(0.2),
                    EqualConv2d(
                        out_channels,
                        sft_out_channels,
                        3,
                        stride=1,
                        padding=1,
                        bias=True,
                        bias_init_val=0,
                    ),
                )
            )

    def forward(
        self, x, return_latents=False, return_rgb=True, randomize_noise=True, **kwargs
    ):
        """Forward function for GFPGANv1.
        Args:
            x (Tensor): Input images.
            return_latents (bool): Whether to return style latents. Default: False.
            return_rgb (bool): Whether return intermediate rgb images. Default: True.
            randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True.
        """
        conditions = []
        unet_skips = []
        out_rgbs = []

        # encoder
        feat = self.conv_body_first(x)
        for i in range(self.log_size - 2):
            feat = self.conv_body_down[i](feat)
            unet_skips.insert(0, feat)

        feat = self.final_conv(feat)

        # style code
        style_code = self.final_linear(feat.view(feat.size(0), -1))
        if self.different_w:
            style_code = style_code.view(style_code.size(0), -1, self.num_style_feat)

        # decode
        for i in range(self.log_size - 2):
            # add unet skip
            feat = feat + unet_skips[i]
            # ResUpLayer
            feat = self.conv_body_up[i](feat)
            # generate scale and shift for SFT layers
            scale = self.condition_scale[i](feat)
            conditions.append(scale.clone())
            shift = self.condition_shift[i](feat)
            conditions.append(shift.clone())
            # generate rgb images
            if return_rgb:
                out_rgbs.append(self.toRGB[i](feat))

        # decoder
        image, _ = self.stylegan_decoder(
            [style_code],
            conditions,
            return_latents=return_latents,
            input_is_latent=self.input_is_latent,
            randomize_noise=randomize_noise,
        )

        return image, out_rgbs


class FacialComponentDiscriminator(nn.Module):
    """Facial component (eyes, mouth, noise) discriminator used in GFPGAN."""

    def __init__(self):
        super(FacialComponentDiscriminator, self).__init__()
        # It now uses a VGG-style architectrue with fixed model size
        self.conv1 = ConvLayer(
            3,
            64,
            3,
            downsample=False,
            resample_kernel=(1, 3, 3, 1),
            bias=True,
            activate=True,
        )
        self.conv2 = ConvLayer(
            64,
            128,
            3,
            downsample=True,
            resample_kernel=(1, 3, 3, 1),
            bias=True,
            activate=True,
        )
        self.conv3 = ConvLayer(
            128,
            128,
            3,
            downsample=False,
            resample_kernel=(1, 3, 3, 1),
            bias=True,
            activate=True,
        )
        self.conv4 = ConvLayer(
            128,
            256,
            3,
            downsample=True,
            resample_kernel=(1, 3, 3, 1),
            bias=True,
            activate=True,
        )
        self.conv5 = ConvLayer(
            256,
            256,
            3,
            downsample=False,
            resample_kernel=(1, 3, 3, 1),
            bias=True,
            activate=True,
        )
        self.final_conv = ConvLayer(256, 1, 3, bias=True, activate=False)

    def forward(self, x, return_feats=False, **kwargs):
        """Forward function for FacialComponentDiscriminator.
        Args:
            x (Tensor): Input images.
            return_feats (bool): Whether to return intermediate features. Default: False.
        """
        feat = self.conv1(x)
        feat = self.conv3(self.conv2(feat))
        rlt_feats = []
        if return_feats:
            rlt_feats.append(feat.clone())
        feat = self.conv5(self.conv4(feat))
        if return_feats:
            rlt_feats.append(feat.clone())
        out = self.final_conv(feat)

        if return_feats:
            return out, rlt_feats
        else:
            return out, None

nodes/impl/pytorch/architecture/face/gfpganv1_arch.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# type: ignore
import math
import random

import torch
from torch import nn
from torch.nn import functional as F

from .stylegan2_clean_arch import StyleGAN2GeneratorClean


class StyleGAN2GeneratorCSFT(StyleGAN2GeneratorClean):
    """StyleGAN2 Generator with SFT modulation (Spatial Feature Transform).
    It is the clean version without custom compiled CUDA extensions used in StyleGAN2.
    Args:
        out_size (int): The spatial size of outputs.
        num_style_feat (int): Channel number of style features. Default: 512.
        num_mlp (int): Layer number of MLP style layers. Default: 8.
        channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
        narrow (float): The narrow ratio for channels. Default: 1.
        sft_half (bool): Whether to apply SFT on half of the input channels. Default: False.
    """

    def __init__(
        self,
        out_size,
        num_style_feat=512,
        num_mlp=8,
        channel_multiplier=2,
        narrow=1,
        sft_half=False,
    ):
        super(StyleGAN2GeneratorCSFT, self).__init__(
            out_size,
            num_style_feat=num_style_feat,
            num_mlp=num_mlp,
            channel_multiplier=channel_multiplier,
            narrow=narrow,
        )
        self.sft_half = sft_half

    def forward(
        self,
        styles,
        conditions,
        input_is_latent=False,
        noise=None,
        randomize_noise=True,
        truncation=1,
        truncation_latent=None,
        inject_index=None,
        return_latents=False,
    ):
        """Forward function for StyleGAN2GeneratorCSFT.
        Args:
            styles (list[Tensor]): Sample codes of styles.
            conditions (list[Tensor]): SFT conditions to generators.
            input_is_latent (bool): Whether input is latent style. Default: False.
            noise (Tensor | None): Input noise or None. Default: None.
            randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True.
            truncation (float): The truncation ratio. Default: 1.
            truncation_latent (Tensor | None): The truncation latent tensor. Default: None.
            inject_index (int | None): The injection index for mixing noise. Default: None.
            return_latents (bool): Whether to return style latents. Default: False.
        """
        # style codes -> latents with Style MLP layer
        if not input_is_latent:
            styles = [self.style_mlp(s) for s in styles]
        # noises
        if noise is None:
            if randomize_noise:
                noise = [None] * self.num_layers  # for each style conv layer
            else:  # use the stored noise
                noise = [
                    getattr(self.noises, f"noise{i}") for i in range(self.num_layers)
                ]
        # style truncation
        if truncation < 1:
            style_truncation = []
            for style in styles:
                style_truncation.append(
                    truncation_latent + truncation * (style - truncation_latent)
                )
            styles = style_truncation
        # get style latents with injection
        if len(styles) == 1:
            inject_index = self.num_latent

            if styles[0].ndim < 3:
                # repeat latent code for all the layers
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            else:  # used for encoder with different latent code for each layer
                latent = styles[0]
        elif len(styles) == 2:  # mixing noises
            if inject_index is None:
                inject_index = random.randint(1, self.num_latent - 1)
            latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            latent2 = (
                styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1)
            )
            latent = torch.cat([latent1, latent2], 1)

        # main generation
        out = self.constant_input(latent.shape[0])
        out = self.style_conv1(out, latent[:, 0], noise=noise[0])
        skip = self.to_rgb1(out, latent[:, 1])

        i = 1
        for conv1, conv2, noise1, noise2, to_rgb in zip(
            self.style_convs[::2],
            self.style_convs[1::2],
            noise[1::2],
            noise[2::2],
            self.to_rgbs,
        ):
            out = conv1(out, latent[:, i], noise=noise1)

            # the conditions may have fewer levels
            if i < len(conditions):
                # SFT part to combine the conditions
                if self.sft_half:  # only apply SFT to half of the channels
                    out_same, out_sft = torch.split(out, int(out.size(1) // 2), dim=1)
                    out_sft = out_sft * conditions[i - 1] + conditions[i]
                    out = torch.cat([out_same, out_sft], dim=1)
                else:  # apply SFT to all the channels
                    out = out * conditions[i - 1] + conditions[i]

            out = conv2(out, latent[:, i + 1], noise=noise2)
            skip = to_rgb(out, latent[:, i + 2], skip)  # feature back to the rgb space
            i += 2

        image = skip

        if return_latents:
            return image, latent
        else:
            return image, None


class ResBlock(nn.Module):
    """Residual block with bilinear upsampling/downsampling.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        mode (str): Upsampling/downsampling mode. Options: down | up. Default: down.
    """

    def __init__(self, in_channels, out_channels, mode="down"):
        super(ResBlock, self).__init__()

        self.conv1 = nn.Conv2d(in_channels, in_channels, 3, 1, 1)
        self.conv2 = nn.Conv2d(in_channels, out_channels, 3, 1, 1)
        self.skip = nn.Conv2d(in_channels, out_channels, 1, bias=False)
        if mode == "down":
            self.scale_factor = 0.5
        elif mode == "up":
            self.scale_factor = 2

    def forward(self, x):
        out = F.leaky_relu_(self.conv1(x), negative_slope=0.2)
        # upsample/downsample
        out = F.interpolate(
            out, scale_factor=self.scale_factor, mode="bilinear", align_corners=False
        )
        out = F.leaky_relu_(self.conv2(out), negative_slope=0.2)
        # skip
        x = F.interpolate(
            x, scale_factor=self.scale_factor, mode="bilinear", align_corners=False
        )
        skip = self.skip(x)
        out = out + skip
        return out


class GFPGANv1Clean(nn.Module):
    """The GFPGAN architecture: Unet + StyleGAN2 decoder with SFT.
    It is the clean version without custom compiled CUDA extensions used in StyleGAN2.
    Ref: GFP-GAN: Towards Real-World Blind Face Restoration with Generative Facial Prior.
    Args:
        out_size (int): The spatial size of outputs.
        num_style_feat (int): Channel number of style features. Default: 512.
        channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
        decoder_load_path (str): The path to the pre-trained decoder model (usually, the StyleGAN2). Default: None.
        fix_decoder (bool): Whether to fix the decoder. Default: True.
        num_mlp (int): Layer number of MLP style layers. Default: 8.
        input_is_latent (bool): Whether input is latent style. Default: False.
        different_w (bool): Whether to use different latent w for different layers. Default: False.
        narrow (float): The narrow ratio for channels. Default: 1.
        sft_half (bool): Whether to apply SFT on half of the input channels. Default: False.
    """

    def __init__(
        self,
        state_dict,
    ):
        super(GFPGANv1Clean, self).__init__()

        out_size = 512
        num_style_feat = 512
        channel_multiplier = 2
        decoder_load_path = None
        fix_decoder = False
        num_mlp = 8
        input_is_latent = True
        different_w = True
        narrow = 1
        sft_half = True

        self.model_arch = "GFPGAN"
        self.sub_type = "Face SR"
        self.scale = 8
        self.in_nc = 3
        self.out_nc = 3
        self.state = state_dict

        self.supports_fp16 = False
        self.supports_bf16 = True
        self.min_size_restriction = 512

        self.input_is_latent = input_is_latent
        self.different_w = different_w
        self.num_style_feat = num_style_feat

        unet_narrow = narrow * 0.5  # by default, use a half of input channels
        channels = {
            "4": int(512 * unet_narrow),
            "8": int(512 * unet_narrow),
            "16": int(512 * unet_narrow),
            "32": int(512 * unet_narrow),
            "64": int(256 * channel_multiplier * unet_narrow),
            "128": int(128 * channel_multiplier * unet_narrow),
            "256": int(64 * channel_multiplier * unet_narrow),
            "512": int(32 * channel_multiplier * unet_narrow),
            "1024": int(16 * channel_multiplier * unet_narrow),
        }

        self.log_size = int(math.log(out_size, 2))
        first_out_size = 2 ** (int(math.log(out_size, 2)))

        self.conv_body_first = nn.Conv2d(3, channels[f"{first_out_size}"], 1)

        # downsample
        in_channels = channels[f"{first_out_size}"]
        self.conv_body_down = nn.ModuleList()
        for i in range(self.log_size, 2, -1):
            out_channels = channels[f"{2**(i - 1)}"]
            self.conv_body_down.append(ResBlock(in_channels, out_channels, mode="down"))
            in_channels = out_channels

        self.final_conv = nn.Conv2d(in_channels, channels["4"], 3, 1, 1)

        # upsample
        in_channels = channels["4"]
        self.conv_body_up = nn.ModuleList()
        for i in range(3, self.log_size + 1):
            out_channels = channels[f"{2**i}"]
            self.conv_body_up.append(ResBlock(in_channels, out_channels, mode="up"))
            in_channels = out_channels

        # to RGB
        self.toRGB = nn.ModuleList()
        for i in range(3, self.log_size + 1):
            self.toRGB.append(nn.Conv2d(channels[f"{2**i}"], 3, 1))

        if different_w:
            linear_out_channel = (int(math.log(out_size, 2)) * 2 - 2) * num_style_feat
        else:
            linear_out_channel = num_style_feat

        self.final_linear = nn.Linear(channels["4"] * 4 * 4, linear_out_channel)

        # the decoder: stylegan2 generator with SFT modulations
        self.stylegan_decoder = StyleGAN2GeneratorCSFT(
            out_size=out_size,
            num_style_feat=num_style_feat,
            num_mlp=num_mlp,
            channel_multiplier=channel_multiplier,
            narrow=narrow,
            sft_half=sft_half,
        )

        # load pre-trained stylegan2 model if necessary
        if decoder_load_path:
            self.stylegan_decoder.load_state_dict(
                torch.load(
                    decoder_load_path, map_location=lambda storage, loc: storage
                )["params_ema"]
            )
        # fix decoder without updating params
        if fix_decoder:
            for _, param in self.stylegan_decoder.named_parameters():
                param.requires_grad = False

        # for SFT modulations (scale and shift)
        self.condition_scale = nn.ModuleList()
        self.condition_shift = nn.ModuleList()
        for i in range(3, self.log_size + 1):
            out_channels = channels[f"{2**i}"]
            if sft_half:
                sft_out_channels = out_channels
            else:
                sft_out_channels = out_channels * 2
            self.condition_scale.append(
                nn.Sequential(
                    nn.Conv2d(out_channels, out_channels, 3, 1, 1),
                    nn.LeakyReLU(0.2, True),
                    nn.Conv2d(out_channels, sft_out_channels, 3, 1, 1),
                )
            )
            self.condition_shift.append(
                nn.Sequential(
                    nn.Conv2d(out_channels, out_channels, 3, 1, 1),
                    nn.LeakyReLU(0.2, True),
                    nn.Conv2d(out_channels, sft_out_channels, 3, 1, 1),
                )
            )
        self.load_state_dict(state_dict)

    def forward(
        self, x, return_latents=False, return_rgb=True, randomize_noise=True, **kwargs
    ):
        """Forward function for GFPGANv1Clean.
        Args:
            x (Tensor): Input images.
            return_latents (bool): Whether to return style latents. Default: False.
            return_rgb (bool): Whether return intermediate rgb images. Default: True.
            randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True.
        """
        conditions = []
        unet_skips = []
        out_rgbs = []

        # encoder
        feat = F.leaky_relu_(self.conv_body_first(x), negative_slope=0.2)
        for i in range(self.log_size - 2):
            feat = self.conv_body_down[i](feat)
            unet_skips.insert(0, feat)
        feat = F.leaky_relu_(self.final_conv(feat), negative_slope=0.2)

        # style code
        style_code = self.final_linear(feat.view(feat.size(0), -1))
        if self.different_w:
            style_code = style_code.view(style_code.size(0), -1, self.num_style_feat)

        # decode
        for i in range(self.log_size - 2):
            # add unet skip
            feat = feat + unet_skips[i]
            # ResUpLayer
            feat = self.conv_body_up[i](feat)
            # generate scale and shift for SFT layers
            scale = self.condition_scale[i](feat)
            conditions.append(scale.clone())
            shift = self.condition_shift[i](feat)
            conditions.append(shift.clone())
            # generate rgb images
            if return_rgb:
                out_rgbs.append(self.toRGB[i](feat))

        # decoder
        image, _ = self.stylegan_decoder(
            [style_code],
            conditions,
            return_latents=return_latents,
            input_is_latent=self.input_is_latent,
            randomize_noise=randomize_noise,
        )

        return image, out_rgbs

nodes/impl/pytorch/architecture/face/gfpganv1_clean_arch.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# type: ignore
"""Modified from https://github.com/wzhouxiff/RestoreFormer
"""
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class VectorQuantizer(nn.Module):
    """
    see https://github.com/MishaLaskin/vqvae/blob/d761a999e2267766400dc646d82d3ac3657771d4/models/quantizer.py
    ____________________________________________
    Discretization bottleneck part of the VQ-VAE.
    Inputs:
    - n_e : number of embeddings
    - e_dim : dimension of embedding
    - beta : commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2
    _____________________________________________
    """

    def __init__(self, n_e, e_dim, beta):
        super(VectorQuantizer, self).__init__()
        self.n_e = n_e
        self.e_dim = e_dim
        self.beta = beta

        self.embedding = nn.Embedding(self.n_e, self.e_dim)
        self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)

    def forward(self, z):
        """
        Inputs the output of the encoder network z and maps it to a discrete
        one-hot vector that is the index of the closest embedding vector e_j
        z (continuous) -> z_q (discrete)
        z.shape = (batch, channel, height, width)
        quantization pipeline:
            1. get encoder input (B,C,H,W)
            2. flatten input to (B*H*W,C)
        """
        # reshape z -> (batch, height, width, channel) and flatten
        z = z.permute(0, 2, 3, 1).contiguous()
        z_flattened = z.view(-1, self.e_dim)
        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z

        d = (
            torch.sum(z_flattened**2, dim=1, keepdim=True)
            + torch.sum(self.embedding.weight**2, dim=1)
            - 2 * torch.matmul(z_flattened, self.embedding.weight.t())
        )

        # could possible replace this here
        # #\start...
        # find closest encodings

        min_value, min_encoding_indices = torch.min(d, dim=1)

        min_encoding_indices = min_encoding_indices.unsqueeze(1)

        min_encodings = torch.zeros(min_encoding_indices.shape[0], self.n_e).to(z)
        min_encodings.scatter_(1, min_encoding_indices, 1)

        # dtype min encodings: torch.float32
        # min_encodings shape: torch.Size([2048, 512])
        # min_encoding_indices.shape: torch.Size([2048, 1])

        # get quantized latent vectors
        z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape)
        # .........\end

        # with:
        # .........\start
        # min_encoding_indices = torch.argmin(d, dim=1)
        # z_q = self.embedding(min_encoding_indices)
        # ......\end......... (TODO)

        # compute loss for embedding
        loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean(
            (z_q - z.detach()) ** 2
        )

        # preserve gradients
        z_q = z + (z_q - z).detach()

        # perplexity

        e_mean = torch.mean(min_encodings, dim=0)
        perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10)))

        # reshape back to match original input shape
        z_q = z_q.permute(0, 3, 1, 2).contiguous()

        return z_q, loss, (perplexity, min_encodings, min_encoding_indices, d)

    def get_codebook_entry(self, indices, shape):
        # shape specifying (batch, height, width, channel)
        # TODO: check for more easy handling with nn.Embedding
        min_encodings = torch.zeros(indices.shape[0], self.n_e).to(indices)
        min_encodings.scatter_(1, indices[:, None], 1)

        # get quantized latent vectors
        z_q = torch.matmul(min_encodings.float(), self.embedding.weight)

        if shape is not None:
            z_q = z_q.view(shape)

            # reshape back to match original input shape
            z_q = z_q.permute(0, 3, 1, 2).contiguous()

        return z_q


# pytorch_diffusion + derived encoder decoder
def nonlinearity(x):
    # swish
    return x * torch.sigmoid(x)


def Normalize(in_channels):
    return torch.nn.GroupNorm(
        num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
    )


class Upsample(nn.Module):
    def __init__(self, in_channels, with_conv):
        super().__init__()
        self.with_conv = with_conv
        if self.with_conv:
            self.conv = torch.nn.Conv2d(
                in_channels, in_channels, kernel_size=3, stride=1, padding=1
            )

    def forward(self, x):
        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
        if self.with_conv:
            x = self.conv(x)
        return x


class Downsample(nn.Module):
    def __init__(self, in_channels, with_conv):
        super().__init__()
        self.with_conv = with_conv
        if self.with_conv:
            # no asymmetric padding in torch conv, must do it ourselves
            self.conv = torch.nn.Conv2d(
                in_channels, in_channels, kernel_size=3, stride=2, padding=0
            )

    def forward(self, x):
        if self.with_conv:
            pad = (0, 1, 0, 1)
            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
            x = self.conv(x)
        else:
            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
        return x


class ResnetBlock(nn.Module):
    def __init__(
        self,
        *,
        in_channels,
        out_channels=None,
        conv_shortcut=False,
        dropout,
        temb_channels=512
    ):
        super().__init__()
        self.in_channels = in_channels
        out_channels = in_channels if out_channels is None else out_channels
        self.out_channels = out_channels
        self.use_conv_shortcut = conv_shortcut

        self.norm1 = Normalize(in_channels)
        self.conv1 = torch.nn.Conv2d(
            in_channels, out_channels, kernel_size=3, stride=1, padding=1
        )
        if temb_channels > 0:
            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
        self.norm2 = Normalize(out_channels)
        self.dropout = torch.nn.Dropout(dropout)
        self.conv2 = torch.nn.Conv2d(
            out_channels, out_channels, kernel_size=3, stride=1, padding=1
        )
        if self.in_channels != self.out_channels:
            if self.use_conv_shortcut:
                self.conv_shortcut = torch.nn.Conv2d(
                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
                )
            else:
                self.nin_shortcut = torch.nn.Conv2d(
                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
                )

    def forward(self, x, temb):
        h = x
        h = self.norm1(h)
        h = nonlinearity(h)
        h = self.conv1(h)

        if temb is not None:
            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]

        h = self.norm2(h)
        h = nonlinearity(h)
        h = self.dropout(h)
        h = self.conv2(h)

        if self.in_channels != self.out_channels:
            if self.use_conv_shortcut:
                x = self.conv_shortcut(x)
            else:
                x = self.nin_shortcut(x)

        return x + h


class MultiHeadAttnBlock(nn.Module):
    def __init__(self, in_channels, head_size=1):
        super().__init__()
        self.in_channels = in_channels
        self.head_size = head_size
        self.att_size = in_channels // head_size
        assert (
            in_channels % head_size == 0
        ), "The size of head should be divided by the number of channels."

        self.norm1 = Normalize(in_channels)
        self.norm2 = Normalize(in_channels)

        self.q = torch.nn.Conv2d(
            in_channels, in_channels, kernel_size=1, stride=1, padding=0
        )
        self.k = torch.nn.Conv2d(
            in_channels, in_channels, kernel_size=1, stride=1, padding=0
        )
        self.v = torch.nn.Conv2d(
            in_channels, in_channels, kernel_size=1, stride=1, padding=0
        )
        self.proj_out = torch.nn.Conv2d(
            in_channels, in_channels, kernel_size=1, stride=1, padding=0
        )
        self.num = 0

    def forward(self, x, y=None):
        h_ = x
        h_ = self.norm1(h_)
        if y is None:
            y = h_
        else:
            y = self.norm2(y)

        q = self.q(y)
        k = self.k(h_)
        v = self.v(h_)

        # compute attention
        b, c, h, w = q.shape
        q = q.reshape(b, self.head_size, self.att_size, h * w)
        q = q.permute(0, 3, 1, 2)  # b, hw, head, att

        k = k.reshape(b, self.head_size, self.att_size, h * w)
        k = k.permute(0, 3, 1, 2)

        v = v.reshape(b, self.head_size, self.att_size, h * w)
        v = v.permute(0, 3, 1, 2)

        q = q.transpose(1, 2)
        v = v.transpose(1, 2)
        k = k.transpose(1, 2).transpose(2, 3)

        scale = int(self.att_size) ** (-0.5)
        q.mul_(scale)
        w_ = torch.matmul(q, k)
        w_ = F.softmax(w_, dim=3)

        w_ = w_.matmul(v)

        w_ = w_.transpose(1, 2).contiguous()  # [b, h*w, head, att]
        w_ = w_.view(b, h, w, -1)
        w_ = w_.permute(0, 3, 1, 2)

        w_ = self.proj_out(w_)

        return x + w_


class MultiHeadEncoder(nn.Module):
    def __init__(
        self,
        ch,
        out_ch,
        ch_mult=(1, 2, 4, 8),
        num_res_blocks=2,
        attn_resolutions=(16,),
        dropout=0.0,
        resamp_with_conv=True,
        in_channels=3,
        resolution=512,
        z_channels=256,
        double_z=True,
        enable_mid=True,
        head_size=1,
        **ignore_kwargs
    ):
        super().__init__()
        self.ch = ch
        self.temb_ch = 0
        self.num_resolutions = len(ch_mult)
        self.num_res_blocks = num_res_blocks
        self.resolution = resolution
        self.in_channels = in_channels
        self.enable_mid = enable_mid

        # downsampling
        self.conv_in = torch.nn.Conv2d(
            in_channels, self.ch, kernel_size=3, stride=1, padding=1
        )

        curr_res = resolution
        in_ch_mult = (1,) + tuple(ch_mult)
        self.down = nn.ModuleList()
        for i_level in range(self.num_resolutions):
            block = nn.ModuleList()
            attn = nn.ModuleList()
            block_in = ch * in_ch_mult[i_level]
            block_out = ch * ch_mult[i_level]
            for _i_block in range(self.num_res_blocks):
                block.append(
                    ResnetBlock(
                        in_channels=block_in,
                        out_channels=block_out,
                        temb_channels=self.temb_ch,
                        dropout=dropout,
                    )
                )
                block_in = block_out
                if curr_res in attn_resolutions:
                    attn.append(MultiHeadAttnBlock(block_in, head_size))
            down = nn.Module()
            down.block = block
            down.attn = attn
            if i_level != self.num_resolutions - 1:
                down.downsample = Downsample(block_in, resamp_with_conv)
                curr_res = curr_res // 2
            self.down.append(down)

        # middle
        if self.enable_mid:
            self.mid = nn.Module()
            self.mid.block_1 = ResnetBlock(
                in_channels=block_in,
                out_channels=block_in,
                temb_channels=self.temb_ch,
                dropout=dropout,
            )
            self.mid.attn_1 = MultiHeadAttnBlock(block_in, head_size)
            self.mid.block_2 = ResnetBlock(
                in_channels=block_in,
                out_channels=block_in,
                temb_channels=self.temb_ch,
                dropout=dropout,
            )

        # end
        self.norm_out = Normalize(block_in)
        self.conv_out = torch.nn.Conv2d(
            block_in,
            2 * z_channels if double_z else z_channels,
            kernel_size=3,
            stride=1,
            padding=1,
        )

    def forward(self, x):
        hs = {}
        # timestep embedding
        temb = None

        # downsampling
        h = self.conv_in(x)
        hs["in"] = h
        for i_level in range(self.num_resolutions):
            for i_block in range(self.num_res_blocks):
                h = self.down[i_level].block[i_block](h, temb)
                if len(self.down[i_level].attn) > 0:
                    h = self.down[i_level].attn[i_block](h)

            if i_level != self.num_resolutions - 1:
                # hs.append(h)
                hs["block_" + str(i_level)] = h
                h = self.down[i_level].downsample(h)

        # middle
        # h = hs[-1]
        if self.enable_mid:
            h = self.mid.block_1(h, temb)
            hs["block_" + str(i_level) + "_atten"] = h
            h = self.mid.attn_1(h)
            h = self.mid.block_2(h, temb)
            hs["mid_atten"] = h

        # end
        h = self.norm_out(h)
        h = nonlinearity(h)
        h = self.conv_out(h)
        # hs.append(h)
        hs["out"] = h

        return hs


class MultiHeadDecoder(nn.Module):
    def __init__(
        self,
        ch,
        out_ch,
        ch_mult=(1, 2, 4, 8),
        num_res_blocks=2,
        attn_resolutions=(16,),
        dropout=0.0,
        resamp_with_conv=True,
        in_channels=3,
        resolution=512,
        z_channels=256,
        give_pre_end=False,
        enable_mid=True,
        head_size=1,
        **ignorekwargs
    ):
        super().__init__()
        self.ch = ch
        self.temb_ch = 0
        self.num_resolutions = len(ch_mult)
        self.num_res_blocks = num_res_blocks
        self.resolution = resolution
        self.in_channels = in_channels
        self.give_pre_end = give_pre_end
        self.enable_mid = enable_mid

        # compute in_ch_mult, block_in and curr_res at lowest res
        block_in = ch * ch_mult[self.num_resolutions - 1]
        curr_res = resolution // 2 ** (self.num_resolutions - 1)
        self.z_shape = (1, z_channels, curr_res, curr_res)
        print(
            "Working with z of shape {} = {} dimensions.".format(
                self.z_shape, np.prod(self.z_shape)
            )
        )

        # z to block_in
        self.conv_in = torch.nn.Conv2d(
            z_channels, block_in, kernel_size=3, stride=1, padding=1
        )

        # middle
        if self.enable_mid:
            self.mid = nn.Module()
            self.mid.block_1 = ResnetBlock(
                in_channels=block_in,
                out_channels=block_in,
                temb_channels=self.temb_ch,
                dropout=dropout,
            )
            self.mid.attn_1 = MultiHeadAttnBlock(block_in, head_size)
            self.mid.block_2 = ResnetBlock(
                in_channels=block_in,
                out_channels=block_in,
                temb_channels=self.temb_ch,
                dropout=dropout,
            )

        # upsampling
        self.up = nn.ModuleList()
        for i_level in reversed(range(self.num_resolutions)):
            block = nn.ModuleList()
            attn = nn.ModuleList()
            block_out = ch * ch_mult[i_level]
            for _i_block in range(self.num_res_blocks + 1):
                block.append(
                    ResnetBlock(
                        in_channels=block_in,
                        out_channels=block_out,
                        temb_channels=self.temb_ch,
                        dropout=dropout,
                    )
                )
                block_in = block_out
                if curr_res in attn_resolutions:
                    attn.append(MultiHeadAttnBlock(block_in, head_size))
            up = nn.Module()
            up.block = block
            up.attn = attn
            if i_level != 0:
                up.upsample = Upsample(block_in, resamp_with_conv)
                curr_res = curr_res * 2
            self.up.insert(0, up)  # prepend to get consistent order

        # end
        self.norm_out = Normalize(block_in)
        self.conv_out = torch.nn.Conv2d(
            block_in, out_ch, kernel_size=3, stride=1, padding=1
        )

    def forward(self, z):
        # assert z.shape[1:] == self.z_shape[1:]
        self.last_z_shape = z.shape

        # timestep embedding
        temb = None

        # z to block_in
        h = self.conv_in(z)

        # middle
        if self.enable_mid:
            h = self.mid.block_1(h, temb)
            h = self.mid.attn_1(h)
            h = self.mid.block_2(h, temb)

        # upsampling
        for i_level in reversed(range(self.num_resolutions)):
            for i_block in range(self.num_res_blocks + 1):
                h = self.up[i_level].block[i_block](h, temb)
                if len(self.up[i_level].attn) > 0:
                    h = self.up[i_level].attn[i_block](h)
            if i_level != 0:
                h = self.up[i_level].upsample(h)

        # end
        if self.give_pre_end:
            return h

        h = self.norm_out(h)
        h = nonlinearity(h)
        h = self.conv_out(h)
        return h


class MultiHeadDecoderTransformer(nn.Module):
    def __init__(
        self,
        ch,
        out_ch,
        ch_mult=(1, 2, 4, 8),
        num_res_blocks=2,
        attn_resolutions=(16,),
        dropout=0.0,
        resamp_with_conv=True,
        in_channels=3,
        resolution=512,
        z_channels=256,
        give_pre_end=False,
        enable_mid=True,
        head_size=1,
        **ignorekwargs
    ):
        super().__init__()
        self.ch = ch
        self.temb_ch = 0
        self.num_resolutions = len(ch_mult)
        self.num_res_blocks = num_res_blocks
        self.resolution = resolution
        self.in_channels = in_channels
        self.give_pre_end = give_pre_end
        self.enable_mid = enable_mid

        # compute in_ch_mult, block_in and curr_res at lowest res
        block_in = ch * ch_mult[self.num_resolutions - 1]
        curr_res = resolution // 2 ** (self.num_resolutions - 1)
        self.z_shape = (1, z_channels, curr_res, curr_res)
        print(
            "Working with z of shape {} = {} dimensions.".format(
                self.z_shape, np.prod(self.z_shape)
            )
        )

        # z to block_in
        self.conv_in = torch.nn.Conv2d(
            z_channels, block_in, kernel_size=3, stride=1, padding=1
        )

        # middle
        if self.enable_mid:
            self.mid = nn.Module()
            self.mid.block_1 = ResnetBlock(
                in_channels=block_in,
                out_channels=block_in,
                temb_channels=self.temb_ch,
                dropout=dropout,
            )
            self.mid.attn_1 = MultiHeadAttnBlock(block_in, head_size)
            self.mid.block_2 = ResnetBlock(
                in_channels=block_in,
                out_channels=block_in,
                temb_channels=self.temb_ch,
                dropout=dropout,
            )

        # upsampling
        self.up = nn.ModuleList()
        for i_level in reversed(range(self.num_resolutions)):
            block = nn.ModuleList()
            attn = nn.ModuleList()
            block_out = ch * ch_mult[i_level]
            for _i_block in range(self.num_res_blocks + 1):
                block.append(
                    ResnetBlock(
                        in_channels=block_in,
                        out_channels=block_out,
                        temb_channels=self.temb_ch,
                        dropout=dropout,
                    )
                )
                block_in = block_out
                if curr_res in attn_resolutions:
                    attn.append(MultiHeadAttnBlock(block_in, head_size))
            up = nn.Module()
            up.block = block
            up.attn = attn
            if i_level != 0:
                up.upsample = Upsample(block_in, resamp_with_conv)
                curr_res = curr_res * 2
            self.up.insert(0, up)  # prepend to get consistent order

        # end
        self.norm_out = Normalize(block_in)
        self.conv_out = torch.nn.Conv2d(
            block_in, out_ch, kernel_size=3, stride=1, padding=1
        )

    def forward(self, z, hs):
        # assert z.shape[1:] == self.z_shape[1:]
        # self.last_z_shape = z.shape

        # timestep embedding
        temb = None

        # z to block_in
        h = self.conv_in(z)

        # middle
        if self.enable_mid:
            h = self.mid.block_1(h, temb)
            h = self.mid.attn_1(h, hs["mid_atten"])
            h = self.mid.block_2(h, temb)

        # upsampling
        for i_level in reversed(range(self.num_resolutions)):
            for i_block in range(self.num_res_blocks + 1):
                h = self.up[i_level].block[i_block](h, temb)
                if len(self.up[i_level].attn) > 0:
                    h = self.up[i_level].attn[i_block](
                        h, hs["block_" + str(i_level) + "_atten"]
                    )
                    # hfeature = h.clone()
            if i_level != 0:
                h = self.up[i_level].upsample(h)

        # end
        if self.give_pre_end:
            return h

        h = self.norm_out(h)
        h = nonlinearity(h)
        h = self.conv_out(h)
        return h


class RestoreFormer(nn.Module):
    def __init__(
        self,
        state_dict,
    ):
        super(RestoreFormer, self).__init__()

        n_embed = 1024
        embed_dim = 256
        ch = 64
        out_ch = 3
        ch_mult = (1, 2, 2, 4, 4, 8)
        num_res_blocks = 2
        attn_resolutions = (16,)
        dropout = 0.0
        in_channels = 3
        resolution = 512
        z_channels = 256
        double_z = False
        enable_mid = True
        fix_decoder = False
        fix_codebook = True
        fix_encoder = False
        head_size = 8

        self.model_arch = "RestoreFormer"
        self.sub_type = "Face SR"
        self.scale = 8
        self.in_nc = 3
        self.out_nc = out_ch
        self.state = state_dict

        self.supports_fp16 = False
        self.supports_bf16 = True
        self.min_size_restriction = 16

        self.encoder = MultiHeadEncoder(
            ch=ch,
            out_ch=out_ch,
            ch_mult=ch_mult,
            num_res_blocks=num_res_blocks,
            attn_resolutions=attn_resolutions,
            dropout=dropout,
            in_channels=in_channels,
            resolution=resolution,
            z_channels=z_channels,
            double_z=double_z,
            enable_mid=enable_mid,
            head_size=head_size,
        )
        self.decoder = MultiHeadDecoderTransformer(
            ch=ch,
            out_ch=out_ch,
            ch_mult=ch_mult,
            num_res_blocks=num_res_blocks,
            attn_resolutions=attn_resolutions,
            dropout=dropout,
            in_channels=in_channels,
            resolution=resolution,
            z_channels=z_channels,
            enable_mid=enable_mid,
            head_size=head_size,
        )

        self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25)

        self.quant_conv = torch.nn.Conv2d(z_channels, embed_dim, 1)
        self.post_quant_conv = torch.nn.Conv2d(embed_dim, z_channels, 1)

        if fix_decoder:
            for _, param in self.decoder.named_parameters():
                param.requires_grad = False
            for _, param in self.post_quant_conv.named_parameters():
                param.requires_grad = False
            for _, param in self.quantize.named_parameters():
                param.requires_grad = False
        elif fix_codebook:
            for _, param in self.quantize.named_parameters():
                param.requires_grad = False

        if fix_encoder:
            for _, param in self.encoder.named_parameters():
                param.requires_grad = False

        self.load_state_dict(state_dict)

    def encode(self, x):
        hs = self.encoder(x)
        h = self.quant_conv(hs["out"])
        quant, emb_loss, info = self.quantize(h)
        return quant, emb_loss, info, hs

    def decode(self, quant, hs):
        quant = self.post_quant_conv(quant)
        dec = self.decoder(quant, hs)

        return dec

    def forward(self, input, **kwargs):
        quant, diff, info, hs = self.encode(input)
        dec = self.decode(quant, hs)

        return dec, None

nodes/impl/pytorch/architecture/face/restoreformer_arch.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# type: ignore
import math
import random

import torch
from torch import nn
from torch.nn import functional as F

from .fused_act import FusedLeakyReLU, fused_leaky_relu
from .upfirdn2d import upfirdn2d


class NormStyleCode(nn.Module):
    def forward(self, x):
        """Normalize the style codes.

        Args:
            x (Tensor): Style codes with shape (b, c).

        Returns:
            Tensor: Normalized tensor.
        """
        return x * torch.rsqrt(torch.mean(x**2, dim=1, keepdim=True) + 1e-8)


def make_resample_kernel(k):
    """Make resampling kernel for UpFirDn.

    Args:
        k (list[int]): A list indicating the 1D resample kernel magnitude.

    Returns:
        Tensor: 2D resampled kernel.
    """
    k = torch.tensor(k, dtype=torch.float32)
    if k.ndim == 1:
        k = k[None, :] * k[:, None]  # to 2D kernel, outer product
    # normalize
    k /= k.sum()
    return k


class UpFirDnUpsample(nn.Module):
    """Upsample, FIR filter, and downsample (upsampole version).

    References:
    1. https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.upfirdn.html  # noqa: E501
    2. http://www.ece.northwestern.edu/local-apps/matlabhelp/toolbox/signal/upfirdn.html  # noqa: E501

    Args:
        resample_kernel (list[int]): A list indicating the 1D resample kernel
            magnitude.
        factor (int): Upsampling scale factor. Default: 2.
    """

    def __init__(self, resample_kernel, factor=2):
        super(UpFirDnUpsample, self).__init__()
        self.kernel = make_resample_kernel(resample_kernel) * (factor**2)
        self.factor = factor

        pad = self.kernel.shape[0] - factor
        self.pad = ((pad + 1) // 2 + factor - 1, pad // 2)

    def forward(self, x):
        out = upfirdn2d(x, self.kernel.type_as(x), up=self.factor, down=1, pad=self.pad)
        return out

    def __repr__(self):
        return f"{self.__class__.__name__}(factor={self.factor})"


class UpFirDnDownsample(nn.Module):
    """Upsample, FIR filter, and downsample (downsampole version).

    Args:
        resample_kernel (list[int]): A list indicating the 1D resample kernel
            magnitude.
        factor (int): Downsampling scale factor. Default: 2.
    """

    def __init__(self, resample_kernel, factor=2):
        super(UpFirDnDownsample, self).__init__()
        self.kernel = make_resample_kernel(resample_kernel)
        self.factor = factor

        pad = self.kernel.shape[0] - factor
        self.pad = ((pad + 1) // 2, pad // 2)

    def forward(self, x):
        out = upfirdn2d(x, self.kernel.type_as(x), up=1, down=self.factor, pad=self.pad)
        return out

    def __repr__(self):
        return f"{self.__class__.__name__}(factor={self.factor})"


class UpFirDnSmooth(nn.Module):
    """Upsample, FIR filter, and downsample (smooth version).

    Args:
        resample_kernel (list[int]): A list indicating the 1D resample kernel
            magnitude.
        upsample_factor (int): Upsampling scale factor. Default: 1.
        downsample_factor (int): Downsampling scale factor. Default: 1.
        kernel_size (int): Kernel size: Default: 1.
    """

    def __init__(
        self, resample_kernel, upsample_factor=1, downsample_factor=1, kernel_size=1
    ):
        super(UpFirDnSmooth, self).__init__()
        self.upsample_factor = upsample_factor
        self.downsample_factor = downsample_factor
        self.kernel = make_resample_kernel(resample_kernel)
        if upsample_factor > 1:
            self.kernel = self.kernel * (upsample_factor**2)

        if upsample_factor > 1:
            pad = (self.kernel.shape[0] - upsample_factor) - (kernel_size - 1)
            self.pad = ((pad + 1) // 2 + upsample_factor - 1, pad // 2 + 1)
        elif downsample_factor > 1:
            pad = (self.kernel.shape[0] - downsample_factor) + (kernel_size - 1)
            self.pad = ((pad + 1) // 2, pad // 2)
        else:
            raise NotImplementedError

    def forward(self, x):
        out = upfirdn2d(x, self.kernel.type_as(x), up=1, down=1, pad=self.pad)
        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(upsample_factor={self.upsample_factor}"
            f", downsample_factor={self.downsample_factor})"
        )


class EqualLinear(nn.Module):
    """Equalized Linear as StyleGAN2.

    Args:
        in_channels (int): Size of each sample.
        out_channels (int): Size of each output sample.
        bias (bool): If set to ``False``, the layer will not learn an additive
            bias. Default: ``True``.
        bias_init_val (float): Bias initialized value. Default: 0.
        lr_mul (float): Learning rate multiplier. Default: 1.
        activation (None | str): The activation after ``linear`` operation.
            Supported: 'fused_lrelu', None. Default: None.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        bias=True,
        bias_init_val=0,
        lr_mul=1,
        activation=None,
    ):
        super(EqualLinear, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.lr_mul = lr_mul
        self.activation = activation
        if self.activation not in ["fused_lrelu", None]:
            raise ValueError(
                f"Wrong activation value in EqualLinear: {activation}"
                "Supported ones are: ['fused_lrelu', None]."
            )
        self.scale = (1 / math.sqrt(in_channels)) * lr_mul

        self.weight = nn.Parameter(torch.randn(out_channels, in_channels).div_(lr_mul))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val))
        else:
            self.register_parameter("bias", None)

    def forward(self, x):
        if self.bias is None:
            bias = None
        else:
            bias = self.bias * self.lr_mul
        if self.activation == "fused_lrelu":
            out = F.linear(x, self.weight * self.scale)
            out = fused_leaky_relu(out, bias)
        else:
            out = F.linear(x, self.weight * self.scale, bias=bias)
        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(in_channels={self.in_channels}, "
            f"out_channels={self.out_channels}, bias={self.bias is not None})"
        )


class ModulatedConv2d(nn.Module):
    """Modulated Conv2d used in StyleGAN2.

    There is no bias in ModulatedConv2d.

    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Size of the convolving kernel.
        num_style_feat (int): Channel number of style features.
        demodulate (bool): Whether to demodulate in the conv layer.
            Default: True.
        sample_mode (str | None): Indicating 'upsample', 'downsample' or None.
            Default: None.
        resample_kernel (list[int]): A list indicating the 1D resample kernel
            magnitude. Default: (1, 3, 3, 1).
        eps (float): A value added to the denominator for numerical stability.
            Default: 1e-8.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        num_style_feat,
        demodulate=True,
        sample_mode=None,
        resample_kernel=(1, 3, 3, 1),
        eps=1e-8,
    ):
        super(ModulatedConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.demodulate = demodulate
        self.sample_mode = sample_mode
        self.eps = eps

        if self.sample_mode == "upsample":
            self.smooth = UpFirDnSmooth(
                resample_kernel,
                upsample_factor=2,
                downsample_factor=1,
                kernel_size=kernel_size,
            )
        elif self.sample_mode == "downsample":
            self.smooth = UpFirDnSmooth(
                resample_kernel,
                upsample_factor=1,
                downsample_factor=2,
                kernel_size=kernel_size,
            )
        elif self.sample_mode is None:
            pass
        else:
            raise ValueError(
                f"Wrong sample mode {self.sample_mode}, "
                "supported ones are ['upsample', 'downsample', None]."
            )

        self.scale = 1 / math.sqrt(in_channels * kernel_size**2)
        # modulation inside each modulated conv
        self.modulation = EqualLinear(
            num_style_feat,
            in_channels,
            bias=True,
            bias_init_val=1,
            lr_mul=1,
            activation=None,
        )

        self.weight = nn.Parameter(
            torch.randn(1, out_channels, in_channels, kernel_size, kernel_size)
        )
        self.padding = kernel_size // 2

    def forward(self, x, style):
        """Forward function.

        Args:
            x (Tensor): Tensor with shape (b, c, h, w).
            style (Tensor): Tensor with shape (b, num_style_feat).

        Returns:
            Tensor: Modulated tensor after convolution.
        """
        b, c, h, w = x.shape  # c = c_in
        # weight modulation
        style = self.modulation(style).view(b, 1, c, 1, 1)
        # self.weight: (1, c_out, c_in, k, k); style: (b, 1, c, 1, 1)
        weight = self.scale * self.weight * style  # (b, c_out, c_in, k, k)

        if self.demodulate:
            demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + self.eps)
            weight = weight * demod.view(b, self.out_channels, 1, 1, 1)

        weight = weight.view(
            b * self.out_channels, c, self.kernel_size, self.kernel_size
        )

        if self.sample_mode == "upsample":
            x = x.view(1, b * c, h, w)
            weight = weight.view(
                b, self.out_channels, c, self.kernel_size, self.kernel_size
            )
            weight = weight.transpose(1, 2).reshape(
                b * c, self.out_channels, self.kernel_size, self.kernel_size
            )
            out = F.conv_transpose2d(x, weight, padding=0, stride=2, groups=b)
            out = out.view(b, self.out_channels, *out.shape[2:4])
            out = self.smooth(out)
        elif self.sample_mode == "downsample":
            x = self.smooth(x)
            x = x.view(1, b * c, *x.shape[2:4])
            out = F.conv2d(x, weight, padding=0, stride=2, groups=b)
            out = out.view(b, self.out_channels, *out.shape[2:4])
        else:
            x = x.view(1, b * c, h, w)
            # weight: (b*c_out, c_in, k, k), groups=b
            out = F.conv2d(x, weight, padding=self.padding, groups=b)
            out = out.view(b, self.out_channels, *out.shape[2:4])

        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(in_channels={self.in_channels}, "
            f"out_channels={self.out_channels}, "
            f"kernel_size={self.kernel_size}, "
            f"demodulate={self.demodulate}, sample_mode={self.sample_mode})"
        )


class StyleConv(nn.Module):
    """Style conv.

    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Size of the convolving kernel.
        num_style_feat (int): Channel number of style features.
        demodulate (bool): Whether demodulate in the conv layer. Default: True.
        sample_mode (str | None): Indicating 'upsample', 'downsample' or None.
            Default: None.
        resample_kernel (list[int]): A list indicating the 1D resample kernel
            magnitude. Default: (1, 3, 3, 1).
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        num_style_feat,
        demodulate=True,
        sample_mode=None,
        resample_kernel=(1, 3, 3, 1),
    ):
        super(StyleConv, self).__init__()
        self.modulated_conv = ModulatedConv2d(
            in_channels,
            out_channels,
            kernel_size,
            num_style_feat,
            demodulate=demodulate,
            sample_mode=sample_mode,
            resample_kernel=resample_kernel,
        )
        self.weight = nn.Parameter(torch.zeros(1))  # for noise injection
        self.activate = FusedLeakyReLU(out_channels)

    def forward(self, x, style, noise=None):
        # modulate
        out = self.modulated_conv(x, style)
        # noise injection
        if noise is None:
            b, _, h, w = out.shape
            noise = out.new_empty(b, 1, h, w).normal_()
        out = out + self.weight * noise
        # activation (with bias)
        out = self.activate(out)
        return out


class ToRGB(nn.Module):
    """To RGB from features.

    Args:
        in_channels (int): Channel number of input.
        num_style_feat (int): Channel number of style features.
        upsample (bool): Whether to upsample. Default: True.
        resample_kernel (list[int]): A list indicating the 1D resample kernel
            magnitude. Default: (1, 3, 3, 1).
    """

    def __init__(
        self, in_channels, num_style_feat, upsample=True, resample_kernel=(1, 3, 3, 1)
    ):
        super(ToRGB, self).__init__()
        if upsample:
            self.upsample = UpFirDnUpsample(resample_kernel, factor=2)
        else:
            self.upsample = None
        self.modulated_conv = ModulatedConv2d(
            in_channels,
            3,
            kernel_size=1,
            num_style_feat=num_style_feat,
            demodulate=False,
            sample_mode=None,
        )
        self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))

    def forward(self, x, style, skip=None):
        """Forward function.

        Args:
            x (Tensor): Feature tensor with shape (b, c, h, w).
            style (Tensor): Tensor with shape (b, num_style_feat).
            skip (Tensor): Base/skip tensor. Default: None.

        Returns:
            Tensor: RGB images.
        """
        out = self.modulated_conv(x, style)
        out = out + self.bias
        if skip is not None:
            if self.upsample:
                skip = self.upsample(skip)
            out = out + skip
        return out


class ConstantInput(nn.Module):
    """Constant input.

    Args:
        num_channel (int): Channel number of constant input.
        size (int): Spatial size of constant input.
    """

    def __init__(self, num_channel, size):
        super(ConstantInput, self).__init__()
        self.weight = nn.Parameter(torch.randn(1, num_channel, size, size))

    def forward(self, batch):
        out = self.weight.repeat(batch, 1, 1, 1)
        return out


class StyleGAN2Generator(nn.Module):
    """StyleGAN2 Generator.

    Args:
        out_size (int): The spatial size of outputs.
        num_style_feat (int): Channel number of style features. Default: 512.
        num_mlp (int): Layer number of MLP style layers. Default: 8.
        channel_multiplier (int): Channel multiplier for large networks of
            StyleGAN2. Default: 2.
        resample_kernel (list[int]): A list indicating the 1D resample kernel
            magnitude. A cross production will be applied to extent 1D resample
            kernel to 2D resample kernel. Default: (1, 3, 3, 1).
        lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01.
        narrow (float): Narrow ratio for channels. Default: 1.0.
    """

    def __init__(
        self,
        out_size,
        num_style_feat=512,
        num_mlp=8,
        channel_multiplier=2,
        resample_kernel=(1, 3, 3, 1),
        lr_mlp=0.01,
        narrow=1,
    ):
        super(StyleGAN2Generator, self).__init__()
        # Style MLP layers
        self.num_style_feat = num_style_feat
        style_mlp_layers = [NormStyleCode()]
        for i in range(num_mlp):
            style_mlp_layers.append(
                EqualLinear(
                    num_style_feat,
                    num_style_feat,
                    bias=True,
                    bias_init_val=0,
                    lr_mul=lr_mlp,
                    activation="fused_lrelu",
                )
            )
        self.style_mlp = nn.Sequential(*style_mlp_layers)

        channels = {
            "4": int(512 * narrow),
            "8": int(512 * narrow),
            "16": int(512 * narrow),
            "32": int(512 * narrow),
            "64": int(256 * channel_multiplier * narrow),
            "128": int(128 * channel_multiplier * narrow),
            "256": int(64 * channel_multiplier * narrow),
            "512": int(32 * channel_multiplier * narrow),
            "1024": int(16 * channel_multiplier * narrow),
        }
        self.channels = channels

        self.constant_input = ConstantInput(channels["4"], size=4)
        self.style_conv1 = StyleConv(
            channels["4"],
            channels["4"],
            kernel_size=3,
            num_style_feat=num_style_feat,
            demodulate=True,
            sample_mode=None,
            resample_kernel=resample_kernel,
        )
        self.to_rgb1 = ToRGB(
            channels["4"],
            num_style_feat,
            upsample=False,
            resample_kernel=resample_kernel,
        )

        self.log_size = int(math.log(out_size, 2))
        self.num_layers = (self.log_size - 2) * 2 + 1
        self.num_latent = self.log_size * 2 - 2

        self.style_convs = nn.ModuleList()
        self.to_rgbs = nn.ModuleList()
        self.noises = nn.Module()

        in_channels = channels["4"]
        # noise
        for layer_idx in range(self.num_layers):
            resolution = 2 ** ((layer_idx + 5) // 2)
            shape = [1, 1, resolution, resolution]
            self.noises.register_buffer(f"noise{layer_idx}", torch.randn(*shape))
        # style convs and to_rgbs
        for i in range(3, self.log_size + 1):
            out_channels = channels[f"{2**i}"]
            self.style_convs.append(
                StyleConv(
                    in_channels,
                    out_channels,
                    kernel_size=3,
                    num_style_feat=num_style_feat,
                    demodulate=True,
                    sample_mode="upsample",
                    resample_kernel=resample_kernel,
                )
            )
            self.style_convs.append(
                StyleConv(
                    out_channels,
                    out_channels,
                    kernel_size=3,
                    num_style_feat=num_style_feat,
                    demodulate=True,
                    sample_mode=None,
                    resample_kernel=resample_kernel,
                )
            )
            self.to_rgbs.append(
                ToRGB(
                    out_channels,
                    num_style_feat,
                    upsample=True,
                    resample_kernel=resample_kernel,
                )
            )
            in_channels = out_channels

    def make_noise(self):
        """Make noise for noise injection."""
        device = self.constant_input.weight.device
        noises = [torch.randn(1, 1, 4, 4, device=device)]

        for i in range(3, self.log_size + 1):
            for _ in range(2):
                noises.append(torch.randn(1, 1, 2**i, 2**i, device=device))

        return noises

    def get_latent(self, x):
        return self.style_mlp(x)

    def mean_latent(self, num_latent):
        latent_in = torch.randn(
            num_latent, self.num_style_feat, device=self.constant_input.weight.device
        )
        latent = self.style_mlp(latent_in).mean(0, keepdim=True)
        return latent

    def forward(
        self,
        styles,
        input_is_latent=False,
        noise=None,
        randomize_noise=True,
        truncation=1,
        truncation_latent=None,
        inject_index=None,
        return_latents=False,
    ):
        """Forward function for StyleGAN2Generator.

        Args:
            styles (list[Tensor]): Sample codes of styles.
            input_is_latent (bool): Whether input is latent style.
                Default: False.
            noise (Tensor | None): Input noise or None. Default: None.
            randomize_noise (bool): Randomize noise, used when 'noise' is
                False. Default: True.
            truncation (float): TODO. Default: 1.
            truncation_latent (Tensor | None): TODO. Default: None.
            inject_index (int | None): The injection index for mixing noise.
                Default: None.
            return_latents (bool): Whether to return style latents.
                Default: False.
        """
        # style codes -> latents with Style MLP layer
        if not input_is_latent:
            styles = [self.style_mlp(s) for s in styles]
        # noises
        if noise is None:
            if randomize_noise:
                noise = [None] * self.num_layers  # for each style conv layer
            else:  # use the stored noise
                noise = [
                    getattr(self.noises, f"noise{i}") for i in range(self.num_layers)
                ]
        # style truncation
        if truncation < 1:
            style_truncation = []
            for style in styles:
                style_truncation.append(
                    truncation_latent + truncation * (style - truncation_latent)
                )
            styles = style_truncation
        # get style latent with injection
        if len(styles) == 1:
            inject_index = self.num_latent

            if styles[0].ndim < 3:
                # repeat latent code for all the layers
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            else:  # used for encoder with different latent code for each layer
                latent = styles[0]
        elif len(styles) == 2:  # mixing noises
            if inject_index is None:
                inject_index = random.randint(1, self.num_latent - 1)
            latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            latent2 = (
                styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1)
            )
            latent = torch.cat([latent1, latent2], 1)

        # main generation
        out = self.constant_input(latent.shape[0])
        out = self.style_conv1(out, latent[:, 0], noise=noise[0])
        skip = self.to_rgb1(out, latent[:, 1])

        i = 1
        for conv1, conv2, noise1, noise2, to_rgb in zip(
            self.style_convs[::2],
            self.style_convs[1::2],
            noise[1::2],
            noise[2::2],
            self.to_rgbs,
        ):
            out = conv1(out, latent[:, i], noise=noise1)
            out = conv2(out, latent[:, i + 1], noise=noise2)
            skip = to_rgb(out, latent[:, i + 2], skip)
            i += 2

        image = skip

        if return_latents:
            return image, latent
        else:
            return image, None


class ScaledLeakyReLU(nn.Module):
    """Scaled LeakyReLU.

    Args:
        negative_slope (float): Negative slope. Default: 0.2.
    """

    def __init__(self, negative_slope=0.2):
        super(ScaledLeakyReLU, self).__init__()
        self.negative_slope = negative_slope

    def forward(self, x):
        out = F.leaky_relu(x, negative_slope=self.negative_slope)
        return out * math.sqrt(2)


class EqualConv2d(nn.Module):
    """Equalized Linear as StyleGAN2.

    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Size of the convolving kernel.
        stride (int): Stride of the convolution. Default: 1
        padding (int): Zero-padding added to both sides of the input.
            Default: 0.
        bias (bool): If ``True``, adds a learnable bias to the output.
            Default: ``True``.
        bias_init_val (float): Bias initialized value. Default: 0.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        padding=0,
        bias=True,
        bias_init_val=0,
    ):
        super(EqualConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.scale = 1 / math.sqrt(in_channels * kernel_size**2)

        self.weight = nn.Parameter(
            torch.randn(out_channels, in_channels, kernel_size, kernel_size)
        )
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val))
        else:
            self.register_parameter("bias", None)

    def forward(self, x):
        out = F.conv2d(
            x,
            self.weight * self.scale,
            bias=self.bias,
            stride=self.stride,
            padding=self.padding,
        )

        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(in_channels={self.in_channels}, "
            f"out_channels={self.out_channels}, "
            f"kernel_size={self.kernel_size},"
            f" stride={self.stride}, padding={self.padding}, "
            f"bias={self.bias is not None})"
        )


class ConvLayer(nn.Sequential):
    """Conv Layer used in StyleGAN2 Discriminator.

    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Kernel size.
        downsample (bool): Whether downsample by a factor of 2.
            Default: False.
        resample_kernel (list[int]): A list indicating the 1D resample
            kernel magnitude. A cross production will be applied to
            extent 1D resample kernel to 2D resample kernel.
            Default: (1, 3, 3, 1).
        bias (bool): Whether with bias. Default: True.
        activate (bool): Whether use activateion. Default: True.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        downsample=False,
        resample_kernel=(1, 3, 3, 1),
        bias=True,
        activate=True,
    ):
        layers = []
        # downsample
        if downsample:
            layers.append(
                UpFirDnSmooth(
                    resample_kernel,
                    upsample_factor=1,
                    downsample_factor=2,
                    kernel_size=kernel_size,
                )
            )
            stride = 2
            self.padding = 0
        else:
            stride = 1
            self.padding = kernel_size // 2
        # conv
        layers.append(
            EqualConv2d(
                in_channels,
                out_channels,
                kernel_size,
                stride=stride,
                padding=self.padding,
                bias=bias and not activate,
            )
        )
        # activation
        if activate:
            if bias:
                layers.append(FusedLeakyReLU(out_channels))
            else:
                layers.append(ScaledLeakyReLU(0.2))

        super(ConvLayer, self).__init__(*layers)


class ResBlock(nn.Module):
    """Residual block used in StyleGAN2 Discriminator.

    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        resample_kernel (list[int]): A list indicating the 1D resample
            kernel magnitude. A cross production will be applied to
            extent 1D resample kernel to 2D resample kernel.
            Default: (1, 3, 3, 1).
    """

    def __init__(self, in_channels, out_channels, resample_kernel=(1, 3, 3, 1)):
        super(ResBlock, self).__init__()

        self.conv1 = ConvLayer(in_channels, in_channels, 3, bias=True, activate=True)
        self.conv2 = ConvLayer(
            in_channels,
            out_channels,
            3,
            downsample=True,
            resample_kernel=resample_kernel,
            bias=True,
            activate=True,
        )
        self.skip = ConvLayer(
            in_channels,
            out_channels,
            1,
            downsample=True,
            resample_kernel=resample_kernel,
            bias=False,
            activate=False,
        )

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        skip = self.skip(x)
        out = (out + skip) / math.sqrt(2)
        return out

nodes/impl/pytorch/architecture/face/stylegan2_arch.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# type: ignore
import math
import random

import torch
from torch import nn
from torch.nn import functional as F

from .fused_act import FusedLeakyReLU, fused_leaky_relu


class NormStyleCode(nn.Module):
    def forward(self, x):
        """Normalize the style codes.
        Args:
            x (Tensor): Style codes with shape (b, c).
        Returns:
            Tensor: Normalized tensor.
        """
        return x * torch.rsqrt(torch.mean(x**2, dim=1, keepdim=True) + 1e-8)


class EqualLinear(nn.Module):
    """Equalized Linear as StyleGAN2.
    Args:
        in_channels (int): Size of each sample.
        out_channels (int): Size of each output sample.
        bias (bool): If set to ``False``, the layer will not learn an additive
            bias. Default: ``True``.
        bias_init_val (float): Bias initialized value. Default: 0.
        lr_mul (float): Learning rate multiplier. Default: 1.
        activation (None | str): The activation after ``linear`` operation.
            Supported: 'fused_lrelu', None. Default: None.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        bias=True,
        bias_init_val=0,
        lr_mul=1,
        activation=None,
    ):
        super(EqualLinear, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.lr_mul = lr_mul
        self.activation = activation
        if self.activation not in ["fused_lrelu", None]:
            raise ValueError(
                f"Wrong activation value in EqualLinear: {activation}"
                "Supported ones are: ['fused_lrelu', None]."
            )
        self.scale = (1 / math.sqrt(in_channels)) * lr_mul

        self.weight = nn.Parameter(torch.randn(out_channels, in_channels).div_(lr_mul))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val))
        else:
            self.register_parameter("bias", None)

    def forward(self, x):
        if self.bias is None:
            bias = None
        else:
            bias = self.bias * self.lr_mul
        if self.activation == "fused_lrelu":
            out = F.linear(x, self.weight * self.scale)
            out = fused_leaky_relu(out, bias)
        else:
            out = F.linear(x, self.weight * self.scale, bias=bias)
        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(in_channels={self.in_channels}, "
            f"out_channels={self.out_channels}, bias={self.bias is not None})"
        )


class ModulatedConv2d(nn.Module):
    """Modulated Conv2d used in StyleGAN2.
    There is no bias in ModulatedConv2d.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Size of the convolving kernel.
        num_style_feat (int): Channel number of style features.
        demodulate (bool): Whether to demodulate in the conv layer.
            Default: True.
        sample_mode (str | None): Indicating 'upsample', 'downsample' or None.
            Default: None.
        eps (float): A value added to the denominator for numerical stability.
            Default: 1e-8.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        num_style_feat,
        demodulate=True,
        sample_mode=None,
        eps=1e-8,
        interpolation_mode="bilinear",
    ):
        super(ModulatedConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.demodulate = demodulate
        self.sample_mode = sample_mode
        self.eps = eps
        self.interpolation_mode = interpolation_mode
        if self.interpolation_mode == "nearest":
            self.align_corners = None
        else:
            self.align_corners = False

        self.scale = 1 / math.sqrt(in_channels * kernel_size**2)
        # modulation inside each modulated conv
        self.modulation = EqualLinear(
            num_style_feat,
            in_channels,
            bias=True,
            bias_init_val=1,
            lr_mul=1,
            activation=None,
        )

        self.weight = nn.Parameter(
            torch.randn(1, out_channels, in_channels, kernel_size, kernel_size)
        )
        self.padding = kernel_size // 2

    def forward(self, x, style):
        """Forward function.
        Args:
            x (Tensor): Tensor with shape (b, c, h, w).
            style (Tensor): Tensor with shape (b, num_style_feat).
        Returns:
            Tensor: Modulated tensor after convolution.
        """
        b, c, h, w = x.shape  # c = c_in
        # weight modulation
        style = self.modulation(style).view(b, 1, c, 1, 1)
        # self.weight: (1, c_out, c_in, k, k); style: (b, 1, c, 1, 1)
        weight = self.scale * self.weight * style  # (b, c_out, c_in, k, k)

        if self.demodulate:
            demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + self.eps)
            weight = weight * demod.view(b, self.out_channels, 1, 1, 1)

        weight = weight.view(
            b * self.out_channels, c, self.kernel_size, self.kernel_size
        )

        if self.sample_mode == "upsample":
            x = F.interpolate(
                x,
                scale_factor=2,
                mode=self.interpolation_mode,
                align_corners=self.align_corners,
            )
        elif self.sample_mode == "downsample":
            x = F.interpolate(
                x,
                scale_factor=0.5,
                mode=self.interpolation_mode,
                align_corners=self.align_corners,
            )

        b, c, h, w = x.shape
        x = x.view(1, b * c, h, w)
        # weight: (b*c_out, c_in, k, k), groups=b
        out = F.conv2d(x, weight, padding=self.padding, groups=b)
        out = out.view(b, self.out_channels, *out.shape[2:4])

        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(in_channels={self.in_channels}, "
            f"out_channels={self.out_channels}, "
            f"kernel_size={self.kernel_size}, "
            f"demodulate={self.demodulate}, sample_mode={self.sample_mode})"
        )


class StyleConv(nn.Module):
    """Style conv.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Size of the convolving kernel.
        num_style_feat (int): Channel number of style features.
        demodulate (bool): Whether demodulate in the conv layer. Default: True.
        sample_mode (str | None): Indicating 'upsample', 'downsample' or None.
            Default: None.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        num_style_feat,
        demodulate=True,
        sample_mode=None,
        interpolation_mode="bilinear",
    ):
        super(StyleConv, self).__init__()
        self.modulated_conv = ModulatedConv2d(
            in_channels,
            out_channels,
            kernel_size,
            num_style_feat,
            demodulate=demodulate,
            sample_mode=sample_mode,
            interpolation_mode=interpolation_mode,
        )
        self.weight = nn.Parameter(torch.zeros(1))  # for noise injection
        self.activate = FusedLeakyReLU(out_channels)

    def forward(self, x, style, noise=None):
        # modulate
        out = self.modulated_conv(x, style)
        # noise injection
        if noise is None:
            b, _, h, w = out.shape
            noise = out.new_empty(b, 1, h, w).normal_()
        out = out + self.weight * noise
        # activation (with bias)
        out = self.activate(out)
        return out


class ToRGB(nn.Module):
    """To RGB from features.
    Args:
        in_channels (int): Channel number of input.
        num_style_feat (int): Channel number of style features.
        upsample (bool): Whether to upsample. Default: True.
    """

    def __init__(
        self, in_channels, num_style_feat, upsample=True, interpolation_mode="bilinear"
    ):
        super(ToRGB, self).__init__()
        self.upsample = upsample
        self.interpolation_mode = interpolation_mode
        if self.interpolation_mode == "nearest":
            self.align_corners = None
        else:
            self.align_corners = False
        self.modulated_conv = ModulatedConv2d(
            in_channels,
            3,
            kernel_size=1,
            num_style_feat=num_style_feat,
            demodulate=False,
            sample_mode=None,
            interpolation_mode=interpolation_mode,
        )
        self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))

    def forward(self, x, style, skip=None):
        """Forward function.
        Args:
            x (Tensor): Feature tensor with shape (b, c, h, w).
            style (Tensor): Tensor with shape (b, num_style_feat).
            skip (Tensor): Base/skip tensor. Default: None.
        Returns:
            Tensor: RGB images.
        """
        out = self.modulated_conv(x, style)
        out = out + self.bias
        if skip is not None:
            if self.upsample:
                skip = F.interpolate(
                    skip,
                    scale_factor=2,
                    mode=self.interpolation_mode,
                    align_corners=self.align_corners,
                )
            out = out + skip
        return out


class ConstantInput(nn.Module):
    """Constant input.
    Args:
        num_channel (int): Channel number of constant input.
        size (int): Spatial size of constant input.
    """

    def __init__(self, num_channel, size):
        super(ConstantInput, self).__init__()
        self.weight = nn.Parameter(torch.randn(1, num_channel, size, size))

    def forward(self, batch):
        out = self.weight.repeat(batch, 1, 1, 1)
        return out


class StyleGAN2GeneratorBilinear(nn.Module):
    """StyleGAN2 Generator.
    Args:
        out_size (int): The spatial size of outputs.
        num_style_feat (int): Channel number of style features. Default: 512.
        num_mlp (int): Layer number of MLP style layers. Default: 8.
        channel_multiplier (int): Channel multiplier for large networks of
            StyleGAN2. Default: 2.
        lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01.
        narrow (float): Narrow ratio for channels. Default: 1.0.
    """

    def __init__(
        self,
        out_size,
        num_style_feat=512,
        num_mlp=8,
        channel_multiplier=2,
        lr_mlp=0.01,
        narrow=1,
        interpolation_mode="bilinear",
    ):
        super(StyleGAN2GeneratorBilinear, self).__init__()
        # Style MLP layers
        self.num_style_feat = num_style_feat
        style_mlp_layers = [NormStyleCode()]
        for i in range(num_mlp):
            style_mlp_layers.append(
                EqualLinear(
                    num_style_feat,
                    num_style_feat,
                    bias=True,
                    bias_init_val=0,
                    lr_mul=lr_mlp,
                    activation="fused_lrelu",
                )
            )
        self.style_mlp = nn.Sequential(*style_mlp_layers)

        channels = {
            "4": int(512 * narrow),
            "8": int(512 * narrow),
            "16": int(512 * narrow),
            "32": int(512 * narrow),
            "64": int(256 * channel_multiplier * narrow),
            "128": int(128 * channel_multiplier * narrow),
            "256": int(64 * channel_multiplier * narrow),
            "512": int(32 * channel_multiplier * narrow),
            "1024": int(16 * channel_multiplier * narrow),
        }
        self.channels = channels

        self.constant_input = ConstantInput(channels["4"], size=4)
        self.style_conv1 = StyleConv(
            channels["4"],
            channels["4"],
            kernel_size=3,
            num_style_feat=num_style_feat,
            demodulate=True,
            sample_mode=None,
            interpolation_mode=interpolation_mode,
        )
        self.to_rgb1 = ToRGB(
            channels["4"],
            num_style_feat,
            upsample=False,
            interpolation_mode=interpolation_mode,
        )

        self.log_size = int(math.log(out_size, 2))
        self.num_layers = (self.log_size - 2) * 2 + 1
        self.num_latent = self.log_size * 2 - 2

        self.style_convs = nn.ModuleList()
        self.to_rgbs = nn.ModuleList()
        self.noises = nn.Module()

        in_channels = channels["4"]
        # noise
        for layer_idx in range(self.num_layers):
            resolution = 2 ** ((layer_idx + 5) // 2)
            shape = [1, 1, resolution, resolution]
            self.noises.register_buffer(f"noise{layer_idx}", torch.randn(*shape))
        # style convs and to_rgbs
        for i in range(3, self.log_size + 1):
            out_channels = channels[f"{2**i}"]
            self.style_convs.append(
                StyleConv(
                    in_channels,
                    out_channels,
                    kernel_size=3,
                    num_style_feat=num_style_feat,
                    demodulate=True,
                    sample_mode="upsample",
                    interpolation_mode=interpolation_mode,
                )
            )
            self.style_convs.append(
                StyleConv(
                    out_channels,
                    out_channels,
                    kernel_size=3,
                    num_style_feat=num_style_feat,
                    demodulate=True,
                    sample_mode=None,
                    interpolation_mode=interpolation_mode,
                )
            )
            self.to_rgbs.append(
                ToRGB(
                    out_channels,
                    num_style_feat,
                    upsample=True,
                    interpolation_mode=interpolation_mode,
                )
            )
            in_channels = out_channels

    def make_noise(self):
        """Make noise for noise injection."""
        device = self.constant_input.weight.device
        noises = [torch.randn(1, 1, 4, 4, device=device)]

        for i in range(3, self.log_size + 1):
            for _ in range(2):
                noises.append(torch.randn(1, 1, 2**i, 2**i, device=device))

        return noises

    def get_latent(self, x):
        return self.style_mlp(x)

    def mean_latent(self, num_latent):
        latent_in = torch.randn(
            num_latent, self.num_style_feat, device=self.constant_input.weight.device
        )
        latent = self.style_mlp(latent_in).mean(0, keepdim=True)
        return latent

    def forward(
        self,
        styles,
        input_is_latent=False,
        noise=None,
        randomize_noise=True,
        truncation=1,
        truncation_latent=None,
        inject_index=None,
        return_latents=False,
    ):
        """Forward function for StyleGAN2Generator.
        Args:
            styles (list[Tensor]): Sample codes of styles.
            input_is_latent (bool): Whether input is latent style.
                Default: False.
            noise (Tensor | None): Input noise or None. Default: None.
            randomize_noise (bool): Randomize noise, used when 'noise' is
                False. Default: True.
            truncation (float): TODO. Default: 1.
            truncation_latent (Tensor | None): TODO. Default: None.
            inject_index (int | None): The injection index for mixing noise.
                Default: None.
            return_latents (bool): Whether to return style latents.
                Default: False.
        """
        # style codes -> latents with Style MLP layer
        if not input_is_latent:
            styles = [self.style_mlp(s) for s in styles]
        # noises
        if noise is None:
            if randomize_noise:
                noise = [None] * self.num_layers  # for each style conv layer
            else:  # use the stored noise
                noise = [
                    getattr(self.noises, f"noise{i}") for i in range(self.num_layers)
                ]
        # style truncation
        if truncation < 1:
            style_truncation = []
            for style in styles:
                style_truncation.append(
                    truncation_latent + truncation * (style - truncation_latent)
                )
            styles = style_truncation
        # get style latent with injection
        if len(styles) == 1:
            inject_index = self.num_latent

            if styles[0].ndim < 3:
                # repeat latent code for all the layers
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            else:  # used for encoder with different latent code for each layer
                latent = styles[0]
        elif len(styles) == 2:  # mixing noises
            if inject_index is None:
                inject_index = random.randint(1, self.num_latent - 1)
            latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            latent2 = (
                styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1)
            )
            latent = torch.cat([latent1, latent2], 1)

        # main generation
        out = self.constant_input(latent.shape[0])
        out = self.style_conv1(out, latent[:, 0], noise=noise[0])
        skip = self.to_rgb1(out, latent[:, 1])

        i = 1
        for conv1, conv2, noise1, noise2, to_rgb in zip(
            self.style_convs[::2],
            self.style_convs[1::2],
            noise[1::2],
            noise[2::2],
            self.to_rgbs,
        ):
            out = conv1(out, latent[:, i], noise=noise1)
            out = conv2(out, latent[:, i + 1], noise=noise2)
            skip = to_rgb(out, latent[:, i + 2], skip)
            i += 2

        image = skip

        if return_latents:
            return image, latent
        else:
            return image, None


class ScaledLeakyReLU(nn.Module):
    """Scaled LeakyReLU.
    Args:
        negative_slope (float): Negative slope. Default: 0.2.
    """

    def __init__(self, negative_slope=0.2):
        super(ScaledLeakyReLU, self).__init__()
        self.negative_slope = negative_slope

    def forward(self, x):
        out = F.leaky_relu(x, negative_slope=self.negative_slope)
        return out * math.sqrt(2)


class EqualConv2d(nn.Module):
    """Equalized Linear as StyleGAN2.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Size of the convolving kernel.
        stride (int): Stride of the convolution. Default: 1
        padding (int): Zero-padding added to both sides of the input.
            Default: 0.
        bias (bool): If ``True``, adds a learnable bias to the output.
            Default: ``True``.
        bias_init_val (float): Bias initialized value. Default: 0.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        padding=0,
        bias=True,
        bias_init_val=0,
    ):
        super(EqualConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.scale = 1 / math.sqrt(in_channels * kernel_size**2)

        self.weight = nn.Parameter(
            torch.randn(out_channels, in_channels, kernel_size, kernel_size)
        )
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val))
        else:
            self.register_parameter("bias", None)

    def forward(self, x):
        out = F.conv2d(
            x,
            self.weight * self.scale,
            bias=self.bias,
            stride=self.stride,
            padding=self.padding,
        )

        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(in_channels={self.in_channels}, "
            f"out_channels={self.out_channels}, "
            f"kernel_size={self.kernel_size},"
            f" stride={self.stride}, padding={self.padding}, "
            f"bias={self.bias is not None})"
        )


class ConvLayer(nn.Sequential):
    """Conv Layer used in StyleGAN2 Discriminator.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Kernel size.
        downsample (bool): Whether downsample by a factor of 2.
            Default: False.
        bias (bool): Whether with bias. Default: True.
        activate (bool): Whether use activateion. Default: True.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        downsample=False,
        bias=True,
        activate=True,
        interpolation_mode="bilinear",
    ):
        layers = []
        self.interpolation_mode = interpolation_mode
        # downsample
        if downsample:
            if self.interpolation_mode == "nearest":
                self.align_corners = None
            else:
                self.align_corners = False

            layers.append(
                torch.nn.Upsample(
                    scale_factor=0.5,
                    mode=interpolation_mode,
                    align_corners=self.align_corners,
                )
            )
        stride = 1
        self.padding = kernel_size // 2
        # conv
        layers.append(
            EqualConv2d(
                in_channels,
                out_channels,
                kernel_size,
                stride=stride,
                padding=self.padding,
                bias=bias and not activate,
            )
        )
        # activation
        if activate:
            if bias:
                layers.append(FusedLeakyReLU(out_channels))
            else:
                layers.append(ScaledLeakyReLU(0.2))

        super(ConvLayer, self).__init__(*layers)


class ResBlock(nn.Module):
    """Residual block used in StyleGAN2 Discriminator.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
    """

    def __init__(self, in_channels, out_channels, interpolation_mode="bilinear"):
        super(ResBlock, self).__init__()

        self.conv1 = ConvLayer(in_channels, in_channels, 3, bias=True, activate=True)
        self.conv2 = ConvLayer(
            in_channels,
            out_channels,
            3,
            downsample=True,
            interpolation_mode=interpolation_mode,
            bias=True,
            activate=True,
        )
        self.skip = ConvLayer(
            in_channels,
            out_channels,
            1,
            downsample=True,
            interpolation_mode=interpolation_mode,
            bias=False,
            activate=False,
        )

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        skip = self.skip(x)
        out = (out + skip) / math.sqrt(2)
        return out

nodes/impl/pytorch/architecture/face/stylegan2_bilinear_arch.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# type: ignore
import math

import torch
from torch import nn
from torch.nn import functional as F
from torch.nn import init
from torch.nn.modules.batchnorm import _BatchNorm


@torch.no_grad()
def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):
    """Initialize network weights.
    Args:
        module_list (list[nn.Module] | nn.Module): Modules to be initialized.
        scale (float): Scale initialized weights, especially for residual
            blocks. Default: 1.
        bias_fill (float): The value to fill bias. Default: 0
        kwargs (dict): Other arguments for initialization function.
    """
    if not isinstance(module_list, list):
        module_list = [module_list]
    for module in module_list:
        for m in module.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, **kwargs)
                m.weight.data *= scale
                if m.bias is not None:
                    m.bias.data.fill_(bias_fill)
            elif isinstance(m, nn.Linear):
                init.kaiming_normal_(m.weight, **kwargs)
                m.weight.data *= scale
                if m.bias is not None:
                    m.bias.data.fill_(bias_fill)
            elif isinstance(m, _BatchNorm):
                init.constant_(m.weight, 1)
                if m.bias is not None:
                    m.bias.data.fill_(bias_fill)


class NormStyleCode(nn.Module):
    def forward(self, x):
        """Normalize the style codes.
        Args:
            x (Tensor): Style codes with shape (b, c).
        Returns:
            Tensor: Normalized tensor.
        """
        return x * torch.rsqrt(torch.mean(x**2, dim=1, keepdim=True) + 1e-8)


class ModulatedConv2d(nn.Module):
    """Modulated Conv2d used in StyleGAN2.
    There is no bias in ModulatedConv2d.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Size of the convolving kernel.
        num_style_feat (int): Channel number of style features.
        demodulate (bool): Whether to demodulate in the conv layer. Default: True.
        sample_mode (str | None): Indicating 'upsample', 'downsample' or None. Default: None.
        eps (float): A value added to the denominator for numerical stability. Default: 1e-8.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        num_style_feat,
        demodulate=True,
        sample_mode=None,
        eps=1e-8,
    ):
        super(ModulatedConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.demodulate = demodulate
        self.sample_mode = sample_mode
        self.eps = eps

        # modulation inside each modulated conv
        self.modulation = nn.Linear(num_style_feat, in_channels, bias=True)
        # initialization
        default_init_weights(
            self.modulation,
            scale=1,
            bias_fill=1,
            a=0,
            mode="fan_in",
            nonlinearity="linear",
        )

        self.weight = nn.Parameter(
            torch.randn(1, out_channels, in_channels, kernel_size, kernel_size)
            / math.sqrt(in_channels * kernel_size**2)
        )
        self.padding = kernel_size // 2

    def forward(self, x, style):
        """Forward function.
        Args:
            x (Tensor): Tensor with shape (b, c, h, w).
            style (Tensor): Tensor with shape (b, num_style_feat).
        Returns:
            Tensor: Modulated tensor after convolution.
        """
        b, c, h, w = x.shape  # c = c_in
        # weight modulation
        style = self.modulation(style).view(b, 1, c, 1, 1)
        # self.weight: (1, c_out, c_in, k, k); style: (b, 1, c, 1, 1)
        weight = self.weight * style  # (b, c_out, c_in, k, k)

        if self.demodulate:
            demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + self.eps)
            weight = weight * demod.view(b, self.out_channels, 1, 1, 1)

        weight = weight.view(
            b * self.out_channels, c, self.kernel_size, self.kernel_size
        )

        # upsample or downsample if necessary
        if self.sample_mode == "upsample":
            x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False)
        elif self.sample_mode == "downsample":
            x = F.interpolate(x, scale_factor=0.5, mode="bilinear", align_corners=False)

        b, c, h, w = x.shape
        x = x.view(1, b * c, h, w)
        # weight: (b*c_out, c_in, k, k), groups=b
        out = F.conv2d(x, weight, padding=self.padding, groups=b)
        out = out.view(b, self.out_channels, *out.shape[2:4])

        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(in_channels={self.in_channels}, out_channels={self.out_channels}, "
            f"kernel_size={self.kernel_size}, demodulate={self.demodulate}, sample_mode={self.sample_mode})"
        )


class StyleConv(nn.Module):
    """Style conv used in StyleGAN2.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Size of the convolving kernel.
        num_style_feat (int): Channel number of style features.
        demodulate (bool): Whether demodulate in the conv layer. Default: True.
        sample_mode (str | None): Indicating 'upsample', 'downsample' or None. Default: None.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        num_style_feat,
        demodulate=True,
        sample_mode=None,
    ):
        super(StyleConv, self).__init__()
        self.modulated_conv = ModulatedConv2d(
            in_channels,
            out_channels,
            kernel_size,
            num_style_feat,
            demodulate=demodulate,
            sample_mode=sample_mode,
        )
        self.weight = nn.Parameter(torch.zeros(1))  # for noise injection
        self.bias = nn.Parameter(torch.zeros(1, out_channels, 1, 1))
        self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, x, style, noise=None):
        # modulate
        out = self.modulated_conv(x, style) * 2**0.5  # for conversion
        # noise injection
        if noise is None:
            b, _, h, w = out.shape
            noise = out.new_empty(b, 1, h, w).normal_()
        out = out + self.weight * noise
        # add bias
        out = out + self.bias
        # activation
        out = self.activate(out)
        return out


class ToRGB(nn.Module):
    """To RGB (image space) from features.
    Args:
        in_channels (int): Channel number of input.
        num_style_feat (int): Channel number of style features.
        upsample (bool): Whether to upsample. Default: True.
    """

    def __init__(self, in_channels, num_style_feat, upsample=True):
        super(ToRGB, self).__init__()
        self.upsample = upsample
        self.modulated_conv = ModulatedConv2d(
            in_channels,
            3,
            kernel_size=1,
            num_style_feat=num_style_feat,
            demodulate=False,
            sample_mode=None,
        )
        self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))

    def forward(self, x, style, skip=None):
        """Forward function.
        Args:
            x (Tensor): Feature tensor with shape (b, c, h, w).
            style (Tensor): Tensor with shape (b, num_style_feat).
            skip (Tensor): Base/skip tensor. Default: None.
        Returns:
            Tensor: RGB images.
        """
        out = self.modulated_conv(x, style)
        out = out + self.bias
        if skip is not None:
            if self.upsample:
                skip = F.interpolate(
                    skip, scale_factor=2, mode="bilinear", align_corners=False
                )
            out = out + skip
        return out


class ConstantInput(nn.Module):
    """Constant input.
    Args:
        num_channel (int): Channel number of constant input.
        size (int): Spatial size of constant input.
    """

    def __init__(self, num_channel, size):
        super(ConstantInput, self).__init__()
        self.weight = nn.Parameter(torch.randn(1, num_channel, size, size))

    def forward(self, batch):
        out = self.weight.repeat(batch, 1, 1, 1)
        return out


class StyleGAN2GeneratorClean(nn.Module):
    """Clean version of StyleGAN2 Generator.
    Args:
        out_size (int): The spatial size of outputs.
        num_style_feat (int): Channel number of style features. Default: 512.
        num_mlp (int): Layer number of MLP style layers. Default: 8.
        channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
        narrow (float): Narrow ratio for channels. Default: 1.0.
    """

    def __init__(
        self, out_size, num_style_feat=512, num_mlp=8, channel_multiplier=2, narrow=1
    ):
        super(StyleGAN2GeneratorClean, self).__init__()
        # Style MLP layers
        self.num_style_feat = num_style_feat
        style_mlp_layers = [NormStyleCode()]
        for i in range(num_mlp):
            style_mlp_layers.extend(
                [
                    nn.Linear(num_style_feat, num_style_feat, bias=True),
                    nn.LeakyReLU(negative_slope=0.2, inplace=True),
                ]
            )
        self.style_mlp = nn.Sequential(*style_mlp_layers)
        # initialization
        default_init_weights(
            self.style_mlp,
            scale=1,
            bias_fill=0,
            a=0.2,
            mode="fan_in",
            nonlinearity="leaky_relu",
        )

        # channel list
        channels = {
            "4": int(512 * narrow),
            "8": int(512 * narrow),
            "16": int(512 * narrow),
            "32": int(512 * narrow),
            "64": int(256 * channel_multiplier * narrow),
            "128": int(128 * channel_multiplier * narrow),
            "256": int(64 * channel_multiplier * narrow),
            "512": int(32 * channel_multiplier * narrow),
            "1024": int(16 * channel_multiplier * narrow),
        }
        self.channels = channels

        self.constant_input = ConstantInput(channels["4"], size=4)
        self.style_conv1 = StyleConv(
            channels["4"],
            channels["4"],
            kernel_size=3,
            num_style_feat=num_style_feat,
            demodulate=True,
            sample_mode=None,
        )
        self.to_rgb1 = ToRGB(channels["4"], num_style_feat, upsample=False)

        self.log_size = int(math.log(out_size, 2))
        self.num_layers = (self.log_size - 2) * 2 + 1
        self.num_latent = self.log_size * 2 - 2

        self.style_convs = nn.ModuleList()
        self.to_rgbs = nn.ModuleList()
        self.noises = nn.Module()

        in_channels = channels["4"]
        # noise
        for layer_idx in range(self.num_layers):
            resolution = 2 ** ((layer_idx + 5) // 2)
            shape = [1, 1, resolution, resolution]
            self.noises.register_buffer(f"noise{layer_idx}", torch.randn(*shape))
        # style convs and to_rgbs
        for i in range(3, self.log_size + 1):
            out_channels = channels[f"{2**i}"]
            self.style_convs.append(
                StyleConv(
                    in_channels,
                    out_channels,
                    kernel_size=3,
                    num_style_feat=num_style_feat,
                    demodulate=True,
                    sample_mode="upsample",
                )
            )
            self.style_convs.append(
                StyleConv(
                    out_channels,
                    out_channels,
                    kernel_size=3,
                    num_style_feat=num_style_feat,
                    demodulate=True,
                    sample_mode=None,
                )
            )
            self.to_rgbs.append(ToRGB(out_channels, num_style_feat, upsample=True))
            in_channels = out_channels

    def make_noise(self):
        """Make noise for noise injection."""
        device = self.constant_input.weight.device
        noises = [torch.randn(1, 1, 4, 4, device=device)]

        for i in range(3, self.log_size + 1):
            for _ in range(2):
                noises.append(torch.randn(1, 1, 2**i, 2**i, device=device))

        return noises

    def get_latent(self, x):
        return self.style_mlp(x)

    def mean_latent(self, num_latent):
        latent_in = torch.randn(
            num_latent, self.num_style_feat, device=self.constant_input.weight.device
        )
        latent = self.style_mlp(latent_in).mean(0, keepdim=True)
        return latent

    def forward(
        self,
        styles,
        input_is_latent=False,
        noise=None,
        randomize_noise=True,
        truncation=1,
        truncation_latent=None,
        inject_index=None,
        return_latents=False,
    ):
        """Forward function for StyleGAN2GeneratorClean.
        Args:
            styles (list[Tensor]): Sample codes of styles.
            input_is_latent (bool): Whether input is latent style. Default: False.
            noise (Tensor | None): Input noise or None. Default: None.
            randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True.
            truncation (float): The truncation ratio. Default: 1.
            truncation_latent (Tensor | None): The truncation latent tensor. Default: None.
            inject_index (int | None): The injection index for mixing noise. Default: None.
            return_latents (bool): Whether to return style latents. Default: False.
        """
        # style codes -> latents with Style MLP layer
        if not input_is_latent:
            styles = [self.style_mlp(s) for s in styles]
        # noises
        if noise is None:
            if randomize_noise:
                noise = [None] * self.num_layers  # for each style conv layer
            else:  # use the stored noise
                noise = [
                    getattr(self.noises, f"noise{i}") for i in range(self.num_layers)
                ]
        # style truncation
        if truncation < 1:
            style_truncation = []
            for style in styles:
                style_truncation.append(
                    truncation_latent + truncation * (style - truncation_latent)
                )
            styles = style_truncation
        # get style latents with injection
        if len(styles) == 1:
            inject_index = self.num_latent

            if styles[0].ndim < 3:
                # repeat latent code for all the layers
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            else:  # used for encoder with different latent code for each layer
                latent = styles[0]
        elif len(styles) == 2:  # mixing noises
            if inject_index is None:
                inject_index = random.randint(1, self.num_latent - 1)
            latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            latent2 = (
                styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1)
            )
            latent = torch.cat([latent1, latent2], 1)

        # main generation
        out = self.constant_input(latent.shape[0])
        out = self.style_conv1(out, latent[:, 0], noise=noise[0])
        skip = self.to_rgb1(out, latent[:, 1])

        i = 1
        for conv1, conv2, noise1, noise2, to_rgb in zip(
            self.style_convs[::2],
            self.style_convs[1::2],
            noise[1::2],
            noise[2::2],
            self.to_rgbs,
        ):
            out = conv1(out, latent[:, i], noise=noise1)
            out = conv2(out, latent[:, i + 1], noise=noise2)
            skip = to_rgb(out, latent[:, i + 2], skip)  # feature back to the rgb space
            i += 2

        image = skip

        if return_latents:
            return image, latent
        else:
            return image, None

nodes/impl/pytorch/architecture/face/stylegan2_clean_arch.py
vladmandic/sd-extension-chainner
# pylint: skip-file
# type: ignore
# modify from https://github.com/rosinality/stylegan2-pytorch/blob/master/op/upfirdn2d.py

import os

import torch
from torch.autograd import Function
from torch.nn import functional as F

upfirdn2d_ext = None


class UpFirDn2dBackward(Function):
    @staticmethod
    def forward(
        ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size
    ):
        up_x, up_y = up
        down_x, down_y = down
        g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad

        grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)

        grad_input = upfirdn2d_ext.upfirdn2d(
            grad_output,
            grad_kernel,
            down_x,
            down_y,
            up_x,
            up_y,
            g_pad_x0,
            g_pad_x1,
            g_pad_y0,
            g_pad_y1,
        )
        grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])

        ctx.save_for_backward(kernel)

        pad_x0, pad_x1, pad_y0, pad_y1 = pad

        ctx.up_x = up_x
        ctx.up_y = up_y
        ctx.down_x = down_x
        ctx.down_y = down_y
        ctx.pad_x0 = pad_x0
        ctx.pad_x1 = pad_x1
        ctx.pad_y0 = pad_y0
        ctx.pad_y1 = pad_y1
        ctx.in_size = in_size
        ctx.out_size = out_size

        return grad_input

    @staticmethod
    def backward(ctx, gradgrad_input):
        (kernel,) = ctx.saved_tensors

        gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)

        gradgrad_out = upfirdn2d_ext.upfirdn2d(
            gradgrad_input,
            kernel,
            ctx.up_x,
            ctx.up_y,
            ctx.down_x,
            ctx.down_y,
            ctx.pad_x0,
            ctx.pad_x1,
            ctx.pad_y0,
            ctx.pad_y1,
        )
        # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0],
        #                                  ctx.out_size[1], ctx.in_size[3])
        gradgrad_out = gradgrad_out.view(
            ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]
        )

        return gradgrad_out, None, None, None, None, None, None, None, None


class UpFirDn2d(Function):
    @staticmethod
    def forward(ctx, input, kernel, up, down, pad):
        up_x, up_y = up
        down_x, down_y = down
        pad_x0, pad_x1, pad_y0, pad_y1 = pad

        kernel_h, kernel_w = kernel.shape
        _, channel, in_h, in_w = input.shape
        ctx.in_size = input.shape

        input = input.reshape(-1, in_h, in_w, 1)

        ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))

        out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
        out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
        ctx.out_size = (out_h, out_w)

        ctx.up = (up_x, up_y)
        ctx.down = (down_x, down_y)
        ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1)

        g_pad_x0 = kernel_w - pad_x0 - 1
        g_pad_y0 = kernel_h - pad_y0 - 1
        g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1
        g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1

        ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)

        out = upfirdn2d_ext.upfirdn2d(
            input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
        )
        # out = out.view(major, out_h, out_w, minor)
        out = out.view(-1, channel, out_h, out_w)

        return out

    @staticmethod
    def backward(ctx, grad_output):
        kernel, grad_kernel = ctx.saved_tensors

        grad_input = UpFirDn2dBackward.apply(
            grad_output,
            kernel,
            grad_kernel,
            ctx.up,
            ctx.down,
            ctx.pad,
            ctx.g_pad,
            ctx.in_size,
            ctx.out_size,
        )

        return grad_input, None, None, None, None


def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
    if input.device.type == "cpu":
        out = upfirdn2d_native(
            input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1]
        )
    else:
        out = UpFirDn2d.apply(
            input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])
        )

    return out


def upfirdn2d_native(
    input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
):
    _, channel, in_h, in_w = input.shape
    input = input.reshape(-1, in_h, in_w, 1)

    _, in_h, in_w, minor = input.shape
    kernel_h, kernel_w = kernel.shape

    out = input.view(-1, in_h, 1, in_w, 1, minor)
    out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
    out = out.view(-1, in_h * up_y, in_w * up_x, minor)

    out = F.pad(
        out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]
    )
    out = out[
        :,
        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
        :,
    ]

    out = out.permute(0, 3, 1, 2)
    out = out.reshape(
        [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
    )
    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
    out = F.conv2d(out, w)
    out = out.reshape(
        -1,
        minor,
        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
    )
    out = out.permute(0, 2, 3, 1)
    out = out[:, ::down_y, ::down_x, :]

    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1

    return out.view(-1, channel, out_h, out_w)

nodes/impl/pytorch/architecture/face/upfirdn2d.py
vladmandic/sd-extension-chainner
"""Code used for this implementation of the MAT helper utils is modified from
lama-cleaner, copyright of Sanster: https://github.com/fenglinglwb/MAT"""

import collections
from itertools import repeat
from typing import Any

import numpy as np
import torch
from torch import conv2d, conv_transpose2d


def normalize_2nd_moment(x, dim=1, eps=1e-8):
    return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt()


class EasyDict(dict):
    """Convenience class that behaves like a dict but allows access with the attribute syntax."""

    def __getattr__(self, name: str) -> Any:
        try:
            return self[name]
        except KeyError:
            raise AttributeError(name)

    def __setattr__(self, name: str, value: Any) -> None:
        self[name] = value

    def __delattr__(self, name: str) -> None:
        del self[name]


activation_funcs = {
    "linear": EasyDict(
        func=lambda x, **_: x,
        def_alpha=0,
        def_gain=1,
        cuda_idx=1,
        ref="",
        has_2nd_grad=False,
    ),
    "relu": EasyDict(
        func=lambda x, **_: torch.nn.functional.relu(x),
        def_alpha=0,
        def_gain=np.sqrt(2),
        cuda_idx=2,
        ref="y",
        has_2nd_grad=False,
    ),
    "lrelu": EasyDict(
        func=lambda x, alpha, **_: torch.nn.functional.leaky_relu(x, alpha),
        def_alpha=0.2,
        def_gain=np.sqrt(2),
        cuda_idx=3,
        ref="y",
        has_2nd_grad=False,
    ),
    "tanh": EasyDict(
        func=lambda x, **_: torch.tanh(x),
        def_alpha=0,
        def_gain=1,
        cuda_idx=4,
        ref="y",
        has_2nd_grad=True,
    ),
    "sigmoid": EasyDict(
        func=lambda x, **_: torch.sigmoid(x),
        def_alpha=0,
        def_gain=1,
        cuda_idx=5,
        ref="y",
        has_2nd_grad=True,
    ),
    "elu": EasyDict(
        func=lambda x, **_: torch.nn.functional.elu(x),
        def_alpha=0,
        def_gain=1,
        cuda_idx=6,
        ref="y",
        has_2nd_grad=True,
    ),
    "selu": EasyDict(
        func=lambda x, **_: torch.nn.functional.selu(x),
        def_alpha=0,
        def_gain=1,
        cuda_idx=7,
        ref="y",
        has_2nd_grad=True,
    ),
    "softplus": EasyDict(
        func=lambda x, **_: torch.nn.functional.softplus(x),
        def_alpha=0,
        def_gain=1,
        cuda_idx=8,
        ref="y",
        has_2nd_grad=True,
    ),
    "swish": EasyDict(
        func=lambda x, **_: torch.sigmoid(x) * x,
        def_alpha=0,
        def_gain=np.sqrt(2),
        cuda_idx=9,
        ref="x",
        has_2nd_grad=True,
    ),
}


def _bias_act_ref(x, b=None, dim=1, act="linear", alpha=None, gain=None, clamp=None):
    """Slow reference implementation of `bias_act()` using standard TensorFlow ops."""
    assert isinstance(x, torch.Tensor)
    assert clamp is None or clamp >= 0
    spec = activation_funcs[act]
    alpha = float(alpha if alpha is not None else spec.def_alpha)
    gain = float(gain if gain is not None else spec.def_gain)
    clamp = float(clamp if clamp is not None else -1)

    # Add bias.
    if b is not None:
        assert isinstance(b, torch.Tensor) and b.ndim == 1
        assert 0 <= dim < x.ndim
        assert b.shape[0] == x.shape[dim]
        x = x + b.reshape([-1 if i == dim else 1 for i in range(x.ndim)]).to(x.device)

    # Evaluate activation function.
    alpha = float(alpha)
    x = spec.func(x, alpha=alpha)

    # Scale by gain.
    gain = float(gain)
    if gain != 1:
        x = x * gain

    # Clamp.
    if clamp >= 0:
        x = x.clamp(-clamp, clamp)  # pylint: disable=invalid-unary-operand-type
    return x


def bias_act(
    x, b=None, dim=1, act="linear", alpha=None, gain=None, clamp=None, impl="ref"
):
    r"""Fused bias and activation function.
    Adds bias `b` to activation tensor `x`, evaluates activation function `act`,
    and scales the result by `gain`. Each of the steps is optional. In most cases,
    the fused op is considerably more efficient than performing the same calculation
    using standard PyTorch ops. It supports first and second order gradients,
    but not third order gradients.
    Args:
        x:      Input activation tensor. Can be of any shape.
        b:      Bias vector, or `None` to disable. Must be a 1D tensor of the same type
                as `x`. The shape must be known, and it must match the dimension of `x`
                corresponding to `dim`.
        dim:    The dimension in `x` corresponding to the elements of `b`.
                The value of `dim` is ignored if `b` is not specified.
        act:    Name of the activation function to evaluate, or `"linear"` to disable.
                Can be e.g. `"relu"`, `"lrelu"`, `"tanh"`, `"sigmoid"`, `"swish"`, etc.
                See `activation_funcs` for a full list. `None` is not allowed.
        alpha:  Shape parameter for the activation function, or `None` to use the default.
        gain:   Scaling factor for the output tensor, or `None` to use default.
                See `activation_funcs` for the default scaling of each activation function.
                If unsure, consider specifying 1.
        clamp:  Clamp the output values to `[-clamp, +clamp]`, or `None` to disable
                the clamping (default).
        impl:   Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
    Returns:
        Tensor of the same shape and datatype as `x`.
    """
    assert isinstance(x, torch.Tensor)
    assert impl in ["ref", "cuda"]
    return _bias_act_ref(
        x=x, b=b, dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp
    )


def setup_filter(
    f,
    device=torch.device("cpu"),
    normalize=True,
    flip_filter=False,
    gain=1,
    separable=None,
):
    r"""Convenience function to setup 2D FIR filter for `upfirdn2d()`.
    Args:
        f:           Torch tensor, numpy array, or python list of the shape
                     `[filter_height, filter_width]` (non-separable),
                     `[filter_taps]` (separable),
                     `[]` (impulse), or
                     `None` (identity).
        device:      Result device (default: cpu).
        normalize:   Normalize the filter so that it retains the magnitude
                     for constant input signal (DC)? (default: True).
        flip_filter: Flip the filter? (default: False).
        gain:        Overall scaling factor for signal magnitude (default: 1).
        separable:   Return a separable filter? (default: select automatically).
    Returns:
        Float32 tensor of the shape
        `[filter_height, filter_width]` (non-separable) or
        `[filter_taps]` (separable).
    """
    # Validate.
    if f is None:
        f = 1
    f = torch.as_tensor(f, dtype=torch.float32)
    assert f.ndim in [0, 1, 2]
    assert f.numel() > 0
    if f.ndim == 0:
        f = f[np.newaxis]

    # Separable?
    if separable is None:
        separable = f.ndim == 1 and f.numel() >= 8
    if f.ndim == 1 and not separable:
        f = f.ger(f)
    assert f.ndim == (1 if separable else 2)

    # Apply normalize, flip, gain, and device.
    if normalize:
        f /= f.sum()
    if flip_filter:
        f = f.flip(list(range(f.ndim)))
    f = f * (gain ** (f.ndim / 2))
    f = f.to(device=device)
    return f


def _get_filter_size(f):
    if f is None:
        return 1, 1

    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
    fw = f.shape[-1]
    fh = f.shape[0]

    fw = int(fw)
    fh = int(fh)
    assert fw >= 1 and fh >= 1
    return fw, fh


def _get_weight_shape(w):
    shape = [int(sz) for sz in w.shape]
    return shape


def _parse_scaling(scaling):
    if isinstance(scaling, int):
        scaling = [scaling, scaling]
    assert isinstance(scaling, (list, tuple))
    assert all(isinstance(x, int) for x in scaling)
    sx, sy = scaling
    assert sx >= 1 and sy >= 1
    return sx, sy


def _parse_padding(padding):
    if isinstance(padding, int):
        padding = [padding, padding]
    assert isinstance(padding, (list, tuple))
    assert all(isinstance(x, int) for x in padding)
    if len(padding) == 2:
        padx, pady = padding
        padding = [padx, padx, pady, pady]
    padx0, padx1, pady0, pady1 = padding
    return padx0, padx1, pady0, pady1


def _ntuple(n):
    def parse(x):
        if isinstance(x, collections.abc.Iterable):
            return x
        return tuple(repeat(x, n))

    return parse


to_2tuple = _ntuple(2)


def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
    """Slow reference implementation of `upfirdn2d()` using standard PyTorch ops."""
    # Validate arguments.
    assert isinstance(x, torch.Tensor) and x.ndim == 4
    if f is None:
        f = torch.ones([1, 1], dtype=torch.float32, device=x.device)
    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
    assert f.dtype == torch.float32 and not f.requires_grad
    batch_size, num_channels, in_height, in_width = x.shape
    # upx, upy = _parse_scaling(up)
    # downx, downy = _parse_scaling(down)

    upx, upy = up, up
    downx, downy = down, down

    # padx0, padx1, pady0, pady1 = _parse_padding(padding)
    padx0, padx1, pady0, pady1 = padding[0], padding[1], padding[2], padding[3]

    # Upsample by inserting zeros.
    x = x.reshape([batch_size, num_channels, in_height, 1, in_width, 1])
    x = torch.nn.functional.pad(x, [0, upx - 1, 0, 0, 0, upy - 1])
    x = x.reshape([batch_size, num_channels, in_height * upy, in_width * upx])

    # Pad or crop.
    x = torch.nn.functional.pad(
        x, [max(padx0, 0), max(padx1, 0), max(pady0, 0), max(pady1, 0)]
    )
    x = x[
        :,
        :,
        max(-pady0, 0) : x.shape[2] - max(-pady1, 0),
        max(-padx0, 0) : x.shape[3] - max(-padx1, 0),
    ]

    # Setup filter.
    f = f * (gain ** (f.ndim / 2))
    f = f.to(x.dtype)
    if not flip_filter:
        f = f.flip(list(range(f.ndim)))

    # Convolve with the filter.
    f = f[np.newaxis, np.newaxis].repeat([num_channels, 1] + [1] * f.ndim)
    if f.ndim == 4:
        x = conv2d(input=x, weight=f, groups=num_channels)
    else:
        x = conv2d(input=x, weight=f.unsqueeze(2), groups=num_channels)
        x = conv2d(input=x, weight=f.unsqueeze(3), groups=num_channels)

    # Downsample by throwing away pixels.
    x = x[:, :, ::downy, ::downx]
    return x


def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, impl="cuda"):
    r"""Pad, upsample, filter, and downsample a batch of 2D images.
    Performs the following sequence of operations for each channel:
    1. Upsample the image by inserting N-1 zeros after each pixel (`up`).
    2. Pad the image with the specified number of zeros on each side (`padding`).
       Negative padding corresponds to cropping the image.
    3. Convolve the image with the specified 2D FIR filter (`f`), shrinking it
       so that the footprint of all output pixels lies within the input image.
    4. Downsample the image by keeping every Nth pixel (`down`).
    This sequence of operations bears close resemblance to scipy.signal.upfirdn().
    The fused op is considerably more efficient than performing the same calculation
    using standard PyTorch ops. It supports gradients of arbitrary order.
    Args:
        x:           Float32/float64/float16 input tensor of the shape
                     `[batch_size, num_channels, in_height, in_width]`.
        f:           Float32 FIR filter of the shape
                     `[filter_height, filter_width]` (non-separable),
                     `[filter_taps]` (separable), or
                     `None` (identity).
        up:          Integer upsampling factor. Can be a single int or a list/tuple
                     `[x, y]` (default: 1).
        down:        Integer downsampling factor. Can be a single int or a list/tuple
                     `[x, y]` (default: 1).
        padding:     Padding with respect to the upsampled image. Can be a single number
                     or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
                     (default: 0).
        flip_filter: False = convolution, True = correlation (default: False).
        gain:        Overall scaling factor for signal magnitude (default: 1).
        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
    Returns:
        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
    """
    # assert isinstance(x, torch.Tensor)
    # assert impl in ['ref', 'cuda']
    return _upfirdn2d_ref(
        x, f, up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain
    )


def upsample2d(x, f, up=2, padding=0, flip_filter=False, gain=1, impl="cuda"):
    r"""Upsample a batch of 2D images using the given 2D FIR filter.
    By default, the result is padded so that its shape is a multiple of the input.
    User-specified padding is applied on top of that, with negative values
    indicating cropping. Pixels outside the image are assumed to be zero.
    Args:
        x:           Float32/float64/float16 input tensor of the shape
                     `[batch_size, num_channels, in_height, in_width]`.
        f:           Float32 FIR filter of the shape
                     `[filter_height, filter_width]` (non-separable),
                     `[filter_taps]` (separable), or
                     `None` (identity).
        up:          Integer upsampling factor. Can be a single int or a list/tuple
                     `[x, y]` (default: 1).
        padding:     Padding with respect to the output. Can be a single number or a
                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
                     (default: 0).
        flip_filter: False = convolution, True = correlation (default: False).
        gain:        Overall scaling factor for signal magnitude (default: 1).
        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
    Returns:
        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
    """
    upx, upy = _parse_scaling(up)
    # upx, upy = up, up
    padx0, padx1, pady0, pady1 = _parse_padding(padding)
    # padx0, padx1, pady0, pady1 = padding, padding, padding, padding
    fw, fh = _get_filter_size(f)
    p = [
        padx0 + (fw + upx - 1) // 2,
        padx1 + (fw - upx) // 2,
        pady0 + (fh + upy - 1) // 2,
        pady1 + (fh - upy) // 2,
    ]
    return upfirdn2d(
        x,
        f,
        up=up,
        padding=p,
        flip_filter=flip_filter,
        gain=gain * upx * upy,
        impl=impl,
    )


class FullyConnectedLayer(torch.nn.Module):
    def __init__(
        self,
        in_features,  # Number of input features.
        out_features,  # Number of output features.
        bias=True,  # Apply additive bias before the activation function?
        activation="linear",  # Activation function: 'relu', 'lrelu', etc.
        lr_multiplier=1,  # Learning rate multiplier.
        bias_init=0,  # Initial value for the additive bias.
    ):
        super().__init__()
        self.weight = torch.nn.Parameter(
            torch.randn([out_features, in_features]) / lr_multiplier
        )
        self.bias = (
            torch.nn.Parameter(torch.full([out_features], np.float32(bias_init)))
            if bias
            else None
        )
        self.activation = activation

        self.weight_gain = lr_multiplier / np.sqrt(in_features)
        self.bias_gain = lr_multiplier

    def forward(self, x):
        w = self.weight * self.weight_gain
        b = self.bias
        if b is not None and self.bias_gain != 1:
            b = b * self.bias_gain

        if self.activation == "linear" and b is not None:
            # out = torch.addmm(b.unsqueeze(0), x, w.t())
            x = x.matmul(w.t().to(x.device))
            out = x + b.reshape(
                [-1 if i == x.ndim - 1 else 1 for i in range(x.ndim)]
            ).to(x.device)
        else:
            x = x.matmul(w.t().to(x.device))
            out = bias_act(x, b, act=self.activation, dim=x.ndim - 1).to(x.device)
        return out


def _conv2d_wrapper(
    x, w, stride=1, padding=0, groups=1, transpose=False, flip_weight=True
):
    """Wrapper for the underlying `conv2d()` and `conv_transpose2d()` implementations."""
    out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w)

    # Flip weight if requested.
    if (
        not flip_weight
    ):  # conv2d() actually performs correlation (flip_weight=True) not convolution (flip_weight=False).
        w = w.flip([2, 3])

    # Workaround performance pitfall in cuDNN 8.0.5, triggered when using
    # 1x1 kernel + memory_format=channels_last + less than 64 channels.
    if (
        kw == 1
        and kh == 1
        and stride == 1
        and padding in [0, [0, 0], (0, 0)]
        and not transpose
    ):
        if x.stride()[1] == 1 and min(out_channels, in_channels_per_group) < 64:
            if out_channels <= 4 and groups == 1:
                in_shape = x.shape
                x = w.squeeze(3).squeeze(2) @ x.reshape(
                    [in_shape[0], in_channels_per_group, -1]
                )
                x = x.reshape([in_shape[0], out_channels, in_shape[2], in_shape[3]])
            else:
                x = x.to(memory_format=torch.contiguous_format)
                w = w.to(memory_format=torch.contiguous_format)
                x = conv2d(x, w, groups=groups)
            return x.to(memory_format=torch.channels_last)

    # Otherwise => execute using conv2d_gradfix.
    op = conv_transpose2d if transpose else conv2d
    return op(x, w, stride=stride, padding=padding, groups=groups)


def conv2d_resample(
    x, w, f=None, up=1, down=1, padding=0, groups=1, flip_weight=True, flip_filter=False
):
    r"""2D convolution with optional up/downsampling.
    Padding is performed only once at the beginning, not between the operations.
    Args:
        x:              Input tensor of shape
                        `[batch_size, in_channels, in_height, in_width]`.
        w:              Weight tensor of shape
                        `[out_channels, in_channels//groups, kernel_height, kernel_width]`.
        f:              Low-pass filter for up/downsampling. Must be prepared beforehand by
                        calling setup_filter(). None = identity (default).
        up:             Integer upsampling factor (default: 1).
        down:           Integer downsampling factor (default: 1).
        padding:        Padding with respect to the upsampled image. Can be a single number
                        or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
                        (default: 0).
        groups:         Split input channels into N groups (default: 1).
        flip_weight:    False = convolution, True = correlation (default: True).
        flip_filter:    False = convolution, True = correlation (default: False).
    Returns:
        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
    """
    # Validate arguments.
    assert isinstance(x, torch.Tensor) and (x.ndim == 4)
    assert isinstance(w, torch.Tensor) and (w.ndim == 4) and (w.dtype == x.dtype)
    assert f is None or (
        isinstance(f, torch.Tensor) and f.ndim in [1, 2] and f.dtype == torch.float32
    )
    assert isinstance(up, int) and (up >= 1)
    assert isinstance(down, int) and (down >= 1)
    # assert isinstance(groups, int) and (groups >= 1), f"!!!!!! groups: {groups} isinstance(groups, int)  {isinstance(groups, int)} {type(groups)}"
    out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w)
    fw, fh = _get_filter_size(f)
    # px0, px1, py0, py1 = _parse_padding(padding)
    px0, px1, py0, py1 = padding, padding, padding, padding

    # Adjust padding to account for up/downsampling.
    if up > 1:
        px0 += (fw + up - 1) // 2
        px1 += (fw - up) // 2
        py0 += (fh + up - 1) // 2
        py1 += (fh - up) // 2
    if down > 1:
        px0 += (fw - down + 1) // 2
        px1 += (fw - down) // 2
        py0 += (fh - down + 1) // 2
        py1 += (fh - down) // 2

    # Fast path: 1x1 convolution with downsampling only => downsample first, then convolve.
    if kw == 1 and kh == 1 and (down > 1 and up == 1):
        x = upfirdn2d(
            x=x, f=f, down=down, padding=[px0, px1, py0, py1], flip_filter=flip_filter
        )
        x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)
        return x

    # Fast path: 1x1 convolution with upsampling only => convolve first, then upsample.
    if kw == 1 and kh == 1 and (up > 1 and down == 1):
        x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)
        x = upfirdn2d(
            x=x,
            f=f,
            up=up,
            padding=[px0, px1, py0, py1],
            gain=up**2,
            flip_filter=flip_filter,
        )
        return x

    # Fast path: downsampling only => use strided convolution.
    if down > 1 and up == 1:
        x = upfirdn2d(x=x, f=f, padding=[px0, px1, py0, py1], flip_filter=flip_filter)
        x = _conv2d_wrapper(
            x=x, w=w, stride=down, groups=groups, flip_weight=flip_weight
        )
        return x

    # Fast path: upsampling with optional downsampling => use transpose strided convolution.
    if up > 1:
        if groups == 1:
            w = w.transpose(0, 1)
        else:
            w = w.reshape(groups, out_channels // groups, in_channels_per_group, kh, kw)
            w = w.transpose(1, 2)
            w = w.reshape(
                groups * in_channels_per_group, out_channels // groups, kh, kw
            )
        px0 -= kw - 1
        px1 -= kw - up
        py0 -= kh - 1
        py1 -= kh - up
        pxt = max(min(-px0, -px1), 0)
        pyt = max(min(-py0, -py1), 0)
        x = _conv2d_wrapper(
            x=x,
            w=w,
            stride=up,
            padding=[pyt, pxt],
            groups=groups,
            transpose=True,
            flip_weight=(not flip_weight),
        )
        x = upfirdn2d(
            x=x,
            f=f,
            padding=[px0 + pxt, px1 + pxt, py0 + pyt, py1 + pyt],
            gain=up**2,
            flip_filter=flip_filter,
        )
        if down > 1:
            x = upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter)
        return x

    # Fast path: no up/downsampling, padding supported by the underlying implementation => use plain conv2d.
    if up == 1 and down == 1:
        if px0 == px1 and py0 == py1 and px0 >= 0 and py0 >= 0:
            return _conv2d_wrapper(
                x=x, w=w, padding=[py0, px0], groups=groups, flip_weight=flip_weight
            )

    # Fallback: Generic reference implementation.
    x = upfirdn2d(
        x=x,
        f=(f if up > 1 else None),
        up=up,
        padding=[px0, px1, py0, py1],
        gain=up**2,
        flip_filter=flip_filter,
    )
    x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)
    if down > 1:
        x = upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter)
    return x


class Conv2dLayer(torch.nn.Module):
    def __init__(
        self,
        in_channels,  # Number of input channels.
        out_channels,  # Number of output channels.
        kernel_size,  # Width and height of the convolution kernel.
        bias=True,  # Apply additive bias before the activation function?
        activation="linear",  # Activation function: 'relu', 'lrelu', etc.
        up=1,  # Integer upsampling factor.
        down=1,  # Integer downsampling factor.
        resample_filter=None,  # Low-pass filter to apply when resampling activations.
        conv_clamp=None,  # Clamp the output to +-X, None = disable clamping.
        channels_last=False,  # Expect the input to have memory_format=channels_last?
        trainable=True,  # Update the weights of this layer during training?
    ):
        if resample_filter is None:
            resample_filter = [1, 3, 3, 1]
        super().__init__()
        self.activation = activation
        self.up = up
        self.down = down
        self.register_buffer("resample_filter", setup_filter(resample_filter))
        self.conv_clamp = conv_clamp
        self.padding = kernel_size // 2
        self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size**2))
        self.act_gain = activation_funcs[activation].def_gain

        memory_format = (
            torch.channels_last if channels_last else torch.contiguous_format
        )
        weight = torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(
            memory_format=memory_format
        )
        bias = torch.zeros([out_channels]) if bias else None
        if trainable:
            self.weight = torch.nn.Parameter(weight)
            self.bias = torch.nn.Parameter(bias) if bias is not None else None
        else:
            self.register_buffer("weight", weight)
            if bias is not None:
                self.register_buffer("bias", bias)
            else:
                self.bias = None

    def forward(self, x, gain=1):
        w = self.weight * self.weight_gain
        x = conv2d_resample(
            x=x,
            w=w,
            f=self.resample_filter,
            up=self.up,
            down=self.down,
            padding=self.padding,
        )

        act_gain = self.act_gain * gain
        act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None
        out = bias_act(
            x, self.bias, act=self.activation, gain=act_gain, clamp=act_clamp
        )
        return out

nodes/impl/pytorch/architecture/mat/utils.py
vladmandic/sd-extension-chainner
""" DropBlock, DropPath

PyTorch implementations of DropBlock and DropPath (Stochastic Depth) regularization layers.

Papers:
DropBlock: A regularization method for convolutional networks (https://arxiv.org/abs/1810.12890)

Deep Networks with Stochastic Depth (https://arxiv.org/abs/1603.09382)

Code:
DropBlock impl inspired by two Tensorflow impl that I liked:
 - https://github.com/tensorflow/tpu/blob/master/models/official/resnet/resnet_model.py#L74
 - https://github.com/clovaai/assembled-cnn/blob/master/nets/blocks.py

Hacked together by / Copyright 2020 Ross Wightman
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


def drop_block_2d(
    x,
    drop_prob: float = 0.1,
    block_size: int = 7,
    gamma_scale: float = 1.0,
    with_noise: bool = False,
    inplace: bool = False,
    batchwise: bool = False,
):
    """DropBlock. See https://arxiv.org/pdf/1810.12890.pdf

    DropBlock with an experimental gaussian noise option. This layer has been tested on a few training
    runs with success, but needs further validation and possibly optimization for lower runtime impact.
    """
    _, C, H, W = x.shape
    total_size = W * H
    clipped_block_size = min(block_size, min(W, H))
    # seed_drop_rate, the gamma parameter
    gamma = (
        gamma_scale
        * drop_prob
        * total_size
        / clipped_block_size**2
        / ((W - block_size + 1) * (H - block_size + 1))
    )

    # Forces the block to be inside the feature map.
    w_i, h_i = torch.meshgrid(
        torch.arange(W).to(x.device), torch.arange(H).to(x.device)
    )
    valid_block = (
        (w_i >= clipped_block_size // 2) & (w_i < W - (clipped_block_size - 1) // 2)
    ) & ((h_i >= clipped_block_size // 2) & (h_i < H - (clipped_block_size - 1) // 2))
    valid_block = torch.reshape(valid_block, (1, 1, H, W)).to(dtype=x.dtype)

    if batchwise:
        # one mask for whole batch, quite a bit faster
        uniform_noise = torch.rand((1, C, H, W), dtype=x.dtype, device=x.device)
    else:
        uniform_noise = torch.rand_like(x)
    block_mask = ((2 - gamma - valid_block + uniform_noise) >= 1).to(dtype=x.dtype)
    block_mask = -F.max_pool2d(
        -block_mask,
        kernel_size=clipped_block_size,  # block_size,
        stride=1,
        padding=clipped_block_size // 2,
    )

    if with_noise:
        normal_noise = (
            torch.randn((1, C, H, W), dtype=x.dtype, device=x.device)
            if batchwise
            else torch.randn_like(x)
        )
        if inplace:
            x.mul_(block_mask).add_(normal_noise * (1 - block_mask))
        else:
            x = x * block_mask + normal_noise * (1 - block_mask)
    else:
        normalize_scale = (
            block_mask.numel() / block_mask.to(dtype=torch.float32).sum().add(1e-7)
        ).to(x.dtype)
        if inplace:
            x.mul_(block_mask * normalize_scale)
        else:
            x = x * block_mask * normalize_scale
    return x


def drop_block_fast_2d(
    x: torch.Tensor,
    drop_prob: float = 0.1,
    block_size: int = 7,
    gamma_scale: float = 1.0,
    with_noise: bool = False,
    inplace: bool = False,
):
    """DropBlock. See https://arxiv.org/pdf/1810.12890.pdf

    DropBlock with an experimental gaussian noise option. Simplied from above without concern for valid
    block mask at edges.
    """
    _, _, H, W = x.shape
    total_size = W * H
    clipped_block_size = min(block_size, min(W, H))
    gamma = (
        gamma_scale
        * drop_prob
        * total_size
        / clipped_block_size**2
        / ((W - block_size + 1) * (H - block_size + 1))
    )

    block_mask = torch.empty_like(x).bernoulli_(gamma)
    block_mask = F.max_pool2d(
        block_mask.to(x.dtype),
        kernel_size=clipped_block_size,
        stride=1,
        padding=clipped_block_size // 2,
    )

    if with_noise:
        normal_noise = torch.empty_like(x).normal_()
        if inplace:
            x.mul_(1.0 - block_mask).add_(normal_noise * block_mask)
        else:
            x = x * (1.0 - block_mask) + normal_noise * block_mask
    else:
        block_mask = 1 - block_mask
        normalize_scale = (
            block_mask.numel() / block_mask.to(dtype=torch.float32).sum().add(1e-6)
        ).to(dtype=x.dtype)
        if inplace:
            x.mul_(block_mask * normalize_scale)
        else:
            x = x * block_mask * normalize_scale
    return x


class DropBlock2d(nn.Module):
    """DropBlock. See https://arxiv.org/pdf/1810.12890.pdf"""

    def __init__(
        self,
        drop_prob: float = 0.1,
        block_size: int = 7,
        gamma_scale: float = 1.0,
        with_noise: bool = False,
        inplace: bool = False,
        batchwise: bool = False,
        fast: bool = True,
    ):
        super(DropBlock2d, self).__init__()
        self.drop_prob = drop_prob
        self.gamma_scale = gamma_scale
        self.block_size = block_size
        self.with_noise = with_noise
        self.inplace = inplace
        self.batchwise = batchwise
        self.fast = fast  # FIXME finish comparisons of fast vs not

    def forward(self, x):
        if not self.training or not self.drop_prob:
            return x
        if self.fast:
            return drop_block_fast_2d(
                x,
                self.drop_prob,
                self.block_size,
                self.gamma_scale,
                self.with_noise,
                self.inplace,
            )
        else:
            return drop_block_2d(
                x,
                self.drop_prob,
                self.block_size,
                self.gamma_scale,
                self.with_noise,
                self.inplace,
                self.batchwise,
            )


def drop_path(
    x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.

    """
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (
        x.ndim - 1
    )  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
    if keep_prob > 0.0 and scale_by_keep:
        random_tensor.div_(keep_prob)
    return x * random_tensor


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""

    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob
        self.scale_by_keep = scale_by_keep

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)

    def extra_repr(self):
        return f"drop_prob={round(self.drop_prob,3):0.3f}"

nodes/impl/pytorch/architecture/timm/drop.py
vladmandic/sd-extension-chainner
""" Layer/Module Helpers
Hacked together by / Copyright 2020 Ross Wightman
"""
import collections.abc
from itertools import repeat


# From PyTorch internals
def _ntuple(n):
    def parse(x):
        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
            return x
        return tuple(repeat(x, n))

    return parse


to_1tuple = _ntuple(1)
to_2tuple = _ntuple(2)
to_3tuple = _ntuple(3)
to_4tuple = _ntuple(4)
to_ntuple = _ntuple


def make_divisible(v, divisor=8, min_value=None, round_limit=0.9):
    min_value = min_value or divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_v < round_limit * v:
        new_v += divisor
    return new_v

nodes/impl/pytorch/architecture/timm/helpers.py
vladmandic/sd-extension-chainner
import math
import warnings

import torch
from torch.nn.init import _calculate_fan_in_and_fan_out


def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn(
            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
            "The distribution of values may be incorrect.",
            stacklevel=2,
        )

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.0))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(
    tensor: torch.Tensor, mean=0.0, std=1.0, a=-2.0, b=2.0
) -> torch.Tensor:
    r"""Fills the input Tensor with values drawn from a truncated
    normal distribution. The values are effectively drawn from the
    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \leq \text{mean} \leq b`.

    NOTE: this impl is similar to the PyTorch trunc_normal_, the bounds [a, b] are
    applied while sampling the normal with mean/std applied, therefore a, b args
    should be adjusted to match the range of mean, std args.

    Args:
        tensor: an n-dimensional `torch.Tensor`
        mean: the mean of the normal distribution
        std: the standard deviation of the normal distribution
        a: the minimum cutoff value
        b: the maximum cutoff value
    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.trunc_normal_(w)
    """
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


def trunc_normal_tf_(
    tensor: torch.Tensor, mean=0.0, std=1.0, a=-2.0, b=2.0
) -> torch.Tensor:
    r"""Fills the input Tensor with values drawn from a truncated
    normal distribution. The values are effectively drawn from the
    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \leq \text{mean} \leq b`.

    NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the
    bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0
    and the result is subsquently scaled and shifted by the mean and std args.

    Args:
        tensor: an n-dimensional `torch.Tensor`
        mean: the mean of the normal distribution
        std: the standard deviation of the normal distribution
        a: the minimum cutoff value
        b: the maximum cutoff value
    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.trunc_normal_(w)
    """
    _no_grad_trunc_normal_(tensor, 0, 1.0, a, b)
    with torch.no_grad():
        tensor.mul_(std).add_(mean)
    return tensor


def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"):
    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
    if mode == "fan_in":
        denom = fan_in
    elif mode == "fan_out":
        denom = fan_out
    elif mode == "fan_avg":
        denom = (fan_in + fan_out) / 2

    variance = scale / denom  # type: ignore

    if distribution == "truncated_normal":
        # constant is stddev of standard normal truncated to (-2, 2)
        trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978)
    elif distribution == "normal":
        tensor.normal_(std=math.sqrt(variance))
    elif distribution == "uniform":
        bound = math.sqrt(3 * variance)
        # pylint: disable=invalid-unary-operand-type
        tensor.uniform_(-bound, bound)
    else:
        raise ValueError(f"invalid distribution {distribution}")


def lecun_normal_(tensor):
    variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal")

nodes/impl/pytorch/architecture/timm/weight_init.py
vladmandic/sd-extension-chainner
from __future__ import annotations
import gc
import numpy as np
import torch
from ..upscale.auto_split import Split, Tiler, auto_split
from .utils import np2tensor, safe_cuda_cache_empty, tensor2np
from .types import PyTorchModel


@torch.inference_mode()
def pytorch_auto_split(
    img: np.ndarray,
    model: PyTorchModel,
    device: torch.device,
    use_fp16: bool,
    tiler: Tiler,
) -> np.ndarray:
    model = model.to(device)
    model = model.half() if use_fp16 else model.float()

    def upscale(img: np.ndarray, _):
        img_tensor = np2tensor(img, change_range=True)

        d_img = None
        try:
            d_img = img_tensor.to(device)
            d_img = d_img.half() if use_fp16 else d_img.float()

            result = model(d_img)
            result = tensor2np(
                result.detach().cpu().detach(),
                change_range=False,
                imtype=np.float32,
            )

            del d_img
            return result
        except RuntimeError as e:
            # Check to see if its actually the CUDA out of memory error
            if "allocate" in str(e) or "CUDA" in str(e):
                # Collect garbage (clear VRAM)
                if d_img is not None:
                    try:
                        d_img.detach().cpu()
                    except:
                        pass
                    del d_img
                gc.collect()
                safe_cuda_cache_empty()
                return Split()
            else:
                # Re-raise the exception if not an OOM error
                raise

    try:
        return auto_split(img, upscale, tiler)
    finally:
        del model
        del device
        gc.collect()
        safe_cuda_cache_empty()

nodes/impl/pytorch/auto_split.py
vladmandic/sd-extension-chainner
from .architecture.DAT import DAT
from .architecture.face.codeformer import CodeFormer
from .architecture.face.gfpganv1_clean_arch import GFPGANv1Clean
from .architecture.face.restoreformer_arch import RestoreFormer
from .architecture.HAT import HAT
from .architecture.LaMa import LaMa
from .architecture.MAT import MAT
from .architecture.OmniSR.OmniSR import OmniSR
from .architecture.RRDB import RRDBNet as ESRGAN
from .architecture.SCUNet import SCUNet
from .architecture.SPSR import SPSRNet as SPSR
from .architecture.SRFormer import SRFormer
from .architecture.SRVGG import SRVGGNetCompact as RealESRGANv2
from .architecture.SwiftSRGAN import Generator as SwiftSRGAN
from .architecture.Swin2SR import Swin2SR
from .architecture.SwinIR import SwinIR
from .types import PyTorchModel


class UnsupportedModel(Exception):
    pass


def load_state_dict(state_dict) -> PyTorchModel:
    # logger.debug(f"Loading state dict into pytorch model arch")

    state_dict_keys = list(state_dict.keys())

    if "params_ema" in state_dict_keys:
        state_dict = state_dict["params_ema"]
    elif "params-ema" in state_dict_keys:
        state_dict = state_dict["params-ema"]
    elif "params" in state_dict_keys:
        state_dict = state_dict["params"]

    state_dict_keys = list(state_dict.keys())
    # SRVGGNet Real-ESRGAN (v2)
    if "body.0.weight" in state_dict_keys and "body.1.weight" in state_dict_keys:
        model = RealESRGANv2(state_dict)
    # SPSR (ESRGAN with lots of extra layers)
    elif "f_HR_conv1.0.weight" in state_dict:
        model = SPSR(state_dict)
    # Swift-SRGAN
    elif (
        "model" in state_dict_keys
        and "initial.cnn.depthwise.weight" in state_dict["model"].keys()
    ):
        model = SwiftSRGAN(state_dict)
    # SwinIR, Swin2SR, SRFormer, HAT
    elif "layers.0.residual_group.blocks.0.norm1.weight" in state_dict_keys:
        if (
            "layers.0.residual_group.blocks.0.attn.aligned_relative_position_index"
            in state_dict_keys
        ):
            model = SRFormer(state_dict)
        elif (
            "layers.0.residual_group.blocks.0.conv_block.cab.0.weight"
            in state_dict_keys
        ):
            model = HAT(state_dict)
        elif "patch_embed.proj.weight" in state_dict_keys:
            model = Swin2SR(state_dict)
        else:
            model = SwinIR(state_dict)
    # GFPGAN
    elif (
        "toRGB.0.weight" in state_dict_keys
        and "stylegan_decoder.style_mlp.1.weight" in state_dict_keys
    ):
        model = GFPGANv1Clean(state_dict)
    # RestoreFormer
    elif (
        "encoder.conv_in.weight" in state_dict_keys
        and "encoder.down.0.block.0.norm1.weight" in state_dict_keys
    ):
        model = RestoreFormer(state_dict)
    elif (
        "encoder.blocks.0.weight" in state_dict_keys
        and "quantize.embedding.weight" in state_dict_keys
    ):
        model = CodeFormer(state_dict)
    # LaMa
    elif (
        "model.model.1.bn_l.running_mean" in state_dict_keys
        or "generator.model.1.bn_l.running_mean" in state_dict_keys
    ):
        model = LaMa(state_dict)
    # MAT
    elif "synthesis.first_stage.conv_first.conv.resample_filter" in state_dict_keys:
        model = MAT(state_dict)
    # Omni-SR
    elif "residual_layer.0.residual_layer.0.layer.0.fn.0.weight" in state_dict_keys:
        model = OmniSR(state_dict)
    # SCUNet
    elif "m_head.0.weight" in state_dict_keys and "m_tail.0.weight" in state_dict_keys:
        model = SCUNet(state_dict)
    # DAT
    elif "layers.0.blocks.2.attn.attn_mask_0" in state_dict_keys:
        model = DAT(state_dict)
    # Regular ESRGAN, "new-arch" ESRGAN, Real-ESRGAN v1
    else:
        try:
            model = ESRGAN(state_dict)
        except:
            # pylint: disable=raise-missing-from
            raise UnsupportedModel
    return model

nodes/impl/pytorch/model_loading.py
vladmandic/sd-extension-chainner
from __future__ import annotations
import gc
import numpy as np
import torch
from ....utils.utils import Region, Size, get_h_w_c
from ...image_op import to_op
from ...upscale.auto_split import Split, auto_split
from ...upscale.grayscale import SplitMode, grayscale_split
from ...upscale.passthrough import passthrough_single_color
from ...upscale.tiler import Tiler
from ..utils import safe_cuda_cache_empty
from .pix_transform import Params, PixTransform


class _PixTiler(Tiler):
    def __init__(self, max_tile_size: int = 2048) -> None:
        self.max_tile_size: int = max_tile_size

    def allow_smaller_tile_size(self) -> bool:
        return False

    def starting_tile_size(self, width: int, height: int, _channels: int) -> Size:
        square = min(width, height, self.max_tile_size)
        return square, square

    def split(self, tile_size: Size) -> Size:
        # half the tile size plus a bit extra to account for overlap
        size = tile_size[0] // 2 + tile_size[0] // 8
        if size < 16:
            raise ValueError("Cannot split any further.")
        return size, size


def _as_3d(img: np.ndarray) -> np.ndarray:
    if img.ndim == 3:
        return img
    return np.expand_dims(img, axis=2)


def pix_transform_auto_split(
    source: np.ndarray,
    guide: np.ndarray,
    device: torch.device,
    params: Params,
    split_mode: SplitMode = SplitMode.LAB,
) -> np.ndarray:
    """
    Automatically splits the source and guide image into segments that can be processed by PixTransform.

    The source and guide image may have any number of channels and any size, also long as the size of the guide image is a whole number (greater than 1) multiple of the size of the source image.
    """

    s_w, s_h, _ = get_h_w_c(source)
    g_w, g_h, _ = get_h_w_c(guide)

    assert (
        g_h > s_h and g_w > s_w
    ), "The guide image mus be larger than the source image."
    assert (
        g_w / s_w == g_w // s_w and g_w / s_w == g_h / s_h
    ), "The size of the guide image must be an integer multiple of the size of the source image (e.g. 2x, 3x, 4x, ...)."

    tiler = _PixTiler()
    scale = g_w // s_w

    def upscale(tile: np.ndarray, region: Region):
        try:
            tile_guide = region.scale(scale).read_from(guide)
            pix_op = to_op(PixTransform)(
                guide_img=np.transpose(_as_3d(tile_guide), (2, 0, 1)),
                device=device,
                params=params,
            )
            # passthrough single colors to speed up alpha channels
            pass_op = to_op(passthrough_single_color)(scale, pix_op)

            return grayscale_split(tile, pass_op, split_mode)
        except RuntimeError as e:
            # Check to see if its actually the CUDA out of memory error
            if "allocate" in str(e) or "CUDA" in str(e):
                # Collect garbage (clear VRAM)
                gc.collect()
                safe_cuda_cache_empty()
                return Split()
            else:
                # Re-raise the exception if not an OOM error
                raise

    try:
        return auto_split(source, upscale, tiler)
    finally:
        del device
        gc.collect()
        safe_cuda_cache_empty()

nodes/impl/pytorch/pix_transform/auto_split.py
vladmandic/sd-extension-chainner
from __future__ import annotations

from dataclasses import dataclass
from typing import Literal, Tuple

import numpy as np
import torch
import torch.optim as optim
import torch.utils.data

from .pix_transform_net import PixTransformNet


@dataclass
class Params:
    spatial_features_input: bool = True
    # spatial color head
    weights_regularizer: Tuple[float, float, float] | None = (0.0001, 0.001, 0.001)
    loss: Literal["mse", "l1"] = "l1"
    lr: float = 0.001
    batch_size: int = 32
    iteration: int = 32 * 1024


def PixTransform(
    source_img: np.ndarray,
    guide_img: np.ndarray,
    device: torch.device,
    params: Params,
) -> np.ndarray:
    if len(guide_img.shape) < 3:
        guide_img = np.expand_dims(guide_img, 0)

    _n_channels, hr_height, hr_width = guide_img.shape

    source_img = source_img.squeeze()
    lr_height, lr_width = source_img.shape

    assert hr_height == hr_width
    assert lr_height == lr_width
    assert hr_height % lr_height == 0

    D = hr_height // lr_height
    M = lr_height
    _N = hr_height

    # normalize guide and source
    guide_img = (
        guide_img - np.mean(guide_img, axis=(1, 2), keepdims=True)
    ) / np.maximum(0.0001, np.std(guide_img, axis=(1, 2), keepdims=True))

    source_img_mean = np.mean(source_img)
    source_img_std = np.std(source_img)
    source_img = (source_img - source_img_mean) / np.maximum(0.0001, source_img_std)

    if params.spatial_features_input:
        x = np.linspace(-0.5, 0.5, hr_width)
        x_grid, y_grid = np.meshgrid(x, x, indexing="ij")

        x_grid = np.expand_dims(x_grid, axis=0)
        y_grid = np.expand_dims(y_grid, axis=0)

        guide_img = np.concatenate([guide_img, x_grid, y_grid], axis=0)

    #### prepare_patches #########################################################################
    # guide_patches is M^2 x C x D x D
    # source_pixels is M^2 x 1

    guide_tensor = torch.from_numpy(guide_img).float().to(device)
    source_tensor = torch.from_numpy(source_img).float().to(device)

    guide_patches = torch.zeros((M * M, guide_tensor.shape[0], D, D)).to(device)
    source_pixels = torch.zeros((M * M, 1)).to(device)
    for i in range(0, M):
        for j in range(0, M):
            guide_patches[j + i * M, :, :, :] = guide_tensor[
                :, i * D : (i + 1) * D, j * D : (j + 1) * D
            ]
            source_pixels[j + i * M] = source_tensor[i : (i + 1), j : (j + 1)]

    train_data = torch.utils.data.TensorDataset(guide_patches, source_pixels)
    train_loader = torch.utils.data.DataLoader(
        train_data, batch_size=params.batch_size, shuffle=True
    )
    ###############################################################################################

    #### setup network ############################################################################
    mynet = (
        PixTransformNet(
            channels_in=guide_tensor.shape[0],
            weights_regularizer=params.weights_regularizer,
        )
        .train()
        .to(device)
    )
    optimizer = optim.Adam(mynet.params_with_regularizer, lr=params.lr)
    if params.loss == "mse":
        myloss = torch.nn.MSELoss()
    elif params.loss == "l1":
        myloss = torch.nn.L1Loss()
    else:
        raise AssertionError("unknown loss!")
    ###############################################################################################

    epochs = params.batch_size * params.iteration // (M * M)
    for _epoch in range(0, epochs):
        for x, y in train_loader:
            optimizer.zero_grad()

            y_pred = mynet(x)
            y_mean_pred = torch.mean(y_pred, dim=[2, 3])

            source_patch_consistency = myloss(y_mean_pred, y)

            source_patch_consistency.backward()
            optimizer.step()

    # compute final prediction, un-normalize, and back to numpy
    mynet.eval()
    predicted_target_img = mynet(guide_tensor.unsqueeze(0)).squeeze()
    predicted_target_img = source_img_mean + source_img_std * predicted_target_img
    predicted_target_img = predicted_target_img.cpu().detach().squeeze().numpy()

    return predicted_target_img

nodes/impl/pytorch/pix_transform/pix_transform.py
vladmandic/sd-extension-chainner
from __future__ import annotations

from typing import Tuple

import torch.nn as nn


class PixTransformNet(nn.Module):
    def __init__(
        self,
        channels_in: int = 5,
        kernel_size: int = 1,
        weights_regularizer: Tuple[float, float, float] | None = None,
    ):
        super(PixTransformNet, self).__init__()

        self.channels_in = channels_in

        self.spatial_net = nn.Sequential(
            nn.Conv2d(2, 32, (1, 1), padding=0),
            nn.ReLU(),
            nn.Conv2d(
                32, 2048, (kernel_size, kernel_size), padding=(kernel_size - 1) // 2
            ),
        )
        self.color_net = nn.Sequential(
            nn.Conv2d(channels_in - 2, 32, (1, 1), padding=0),
            nn.ReLU(),
            nn.Conv2d(
                32, 2048, (kernel_size, kernel_size), padding=(kernel_size - 1) // 2
            ),
        )
        self.head_net = nn.Sequential(
            nn.ReLU(),
            nn.Conv2d(
                2048, 32, (kernel_size, kernel_size), padding=(kernel_size - 1) // 2
            ),
            nn.ReLU(),
            nn.Conv2d(32, 1, (1, 1), padding=0),
        )

        if weights_regularizer is None:
            reg_spatial = 0.0001
            reg_color = 0.001
            reg_head = 0.0001
        else:
            reg_spatial = weights_regularizer[0]
            reg_color = weights_regularizer[1]
            reg_head = weights_regularizer[2]

        self.params_with_regularizer = []
        self.params_with_regularizer += [
            {"params": self.spatial_net.parameters(), "weight_decay": reg_spatial}
        ]
        self.params_with_regularizer += [
            {"params": self.color_net.parameters(), "weight_decay": reg_color}
        ]
        self.params_with_regularizer += [
            {"params": self.head_net.parameters(), "weight_decay": reg_head}
        ]

    def forward(self, input_):
        input_spatial = input_[:, self.channels_in - 2 :, :, :]
        input_color = input_[:, 0 : self.channels_in - 2, :, :]

        merged_features = self.spatial_net(input_spatial) + self.color_net(input_color)

        return self.head_net(merged_features)

nodes/impl/pytorch/pix_transform/pix_transform_net.py
vladmandic/sd-extension-chainner
from typing import Union

from .architecture.DAT import DAT
from .architecture.face.codeformer import CodeFormer
from .architecture.face.gfpganv1_clean_arch import GFPGANv1Clean
from .architecture.face.restoreformer_arch import RestoreFormer
from .architecture.HAT import HAT
from .architecture.LaMa import LaMa
from .architecture.MAT import MAT
from .architecture.OmniSR.OmniSR import OmniSR
from .architecture.RRDB import RRDBNet as ESRGAN
from .architecture.SCUNet import SCUNet
from .architecture.SPSR import SPSRNet as SPSR
from .architecture.SRFormer import SRFormer
from .architecture.SRVGG import SRVGGNetCompact as RealESRGANv2
from .architecture.SwiftSRGAN import Generator as SwiftSRGAN
from .architecture.Swin2SR import Swin2SR
from .architecture.SwinIR import SwinIR

PyTorchSRModels = (
    RealESRGANv2,
    SPSR,
    SwiftSRGAN,
    ESRGAN,
    SwinIR,
    Swin2SR,
    HAT,
    OmniSR,
    SCUNet,
    SRFormer,
    DAT,
)
PyTorchSRModel = Union[
    RealESRGANv2,
    SPSR,
    SwiftSRGAN,
    ESRGAN,
    SwinIR,
    Swin2SR,
    HAT,
    OmniSR,
    SCUNet,
    SRFormer,
    DAT,
]


def is_pytorch_sr_model(model: object):
    return isinstance(model, PyTorchSRModels)


PyTorchFaceModels = (GFPGANv1Clean, RestoreFormer, CodeFormer)
PyTorchFaceModel = Union[GFPGANv1Clean, RestoreFormer, CodeFormer]


def is_pytorch_face_model(model: object):
    return isinstance(model, PyTorchFaceModels)


PyTorchInpaintModels = (LaMa, MAT)
PyTorchInpaintModel = Union[LaMa, MAT]


def is_pytorch_inpaint_model(model: object):
    return isinstance(model, PyTorchInpaintModels)


PyTorchModels = (*PyTorchSRModels, *PyTorchFaceModels, *PyTorchInpaintModels)
PyTorchModel = Union[PyTorchSRModel, PyTorchFaceModel, PyTorchInpaintModel]


def is_pytorch_model(model: object):
    return isinstance(model, PyTorchModels)

nodes/impl/pytorch/types.py
vladmandic/sd-extension-chainner
from __future__ import annotations

from typing import Type

import numpy as np
import torch
from torch import Tensor

from ..image_utils import as_3d
from ..onnx.np_tensor_utils import MAX_VALUES_BY_DTYPE, np_denorm


def bgr_to_rgb(image: Tensor) -> Tensor:
    # flip image channels
    # https://github.com/pytorch/pytorch/issues/229
    out: Tensor = image.flip(-3)
    # RGB to BGR #may be faster:
    # out: Tensor = image[[2, 1, 0], :, :]
    return out


def rgb_to_bgr(image: Tensor) -> Tensor:
    # same operation as bgr_to_rgb(), flip image channels
    return bgr_to_rgb(image)


def bgra_to_rgba(image: Tensor) -> Tensor:
    out: Tensor = image[[2, 1, 0, 3], :, :]
    return out


def rgba_to_bgra(image: Tensor) -> Tensor:
    # same operation as bgra_to_rgba(), flip image channels
    return bgra_to_rgba(image)


def norm(x: Tensor):
    """Normalize (z-norm) from [0,1] range to [-1,1]"""
    out = (x - 0.5) * 2.0
    return out.clamp(-1, 1)


def np2tensor(
    img: np.ndarray,
    bgr2rgb=True,
    data_range=1.0,  # pylint: disable=unused-argument
    normalize=False,
    change_range=True,
    add_batch=True,
) -> Tensor:
    """Converts a numpy image array into a Tensor array.
    Parameters:
        img (numpy array): the input image numpy array
        add_batch (bool): choose if new tensor needs batch dimension added
    """

    # check how many channels the image has, then condition. ie. RGB, RGBA, Gray
    # if bgr2rgb:
    #     img = img[
    #         :, :, [2, 1, 0]
    #     ]  # BGR to RGB -> in numpy, if using OpenCV, else not needed. Only if image has colors.
    if change_range:
        dtype = img.dtype
        maxval = MAX_VALUES_BY_DTYPE.get(dtype.name, 1.0)
        t_dtype = np.dtype("float32")
        img = img.astype(t_dtype) / maxval  # ie: uint8 = /255
    # "HWC to CHW" and "numpy to tensor"
    tensor = torch.from_numpy(
        np.ascontiguousarray(np.transpose(as_3d(img), (2, 0, 1)))
    ).float()
    if bgr2rgb:
        # BGR to RGB -> in tensor, if using OpenCV, else not needed. Only if image has colors.)
        if tensor.shape[0] % 3 == 0:
            # RGB or MultixRGB (3xRGB, 5xRGB, etc. For video tensors.)
            tensor = bgr_to_rgb(tensor)
        elif tensor.shape[0] == 4:
            # RGBA
            tensor = bgra_to_rgba(tensor)
    if add_batch:
        # Add fake batch dimension = 1 . squeeze() will remove the dimensions of size 1
        tensor.unsqueeze_(0)
    if normalize:
        tensor = norm(tensor)
    return tensor


def tensor2np(
    img: Tensor,
    rgb2bgr=True,
    remove_batch=True,
    data_range=255,
    denormalize=False,
    change_range=True,
    imtype: Type = np.uint8,
) -> np.ndarray:
    """Converts a Tensor array into a numpy image array.
    Parameters:
        img (tensor): the input image tensor array
            4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order
        remove_batch (bool): choose if tensor of shape BCHW needs to be squeezed
        denormalize (bool): Used to denormalize from [-1,1] range back to [0,1]
        imtype (type): the desired type of the converted numpy array (np.uint8
            default)
    Output:
        img (np array): 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)
    """
    n_dim = img.dim()

    # TODO: Check: could denormalize here in tensor form instead, but end result is the same

    img = img.float().cpu()

    img_np: np.ndarray

    if n_dim in (4, 3):
        # if n_dim == 4, has to convert to 3 dimensions
        if n_dim == 4 and remove_batch:
            # remove a fake batch dimension
            img = img.squeeze(dim=0)

        if img.shape[0] == 3 and rgb2bgr:  # RGB
            # RGB to BGR -> in tensor, if using OpenCV, else not needed. Only if image has colors.
            img_np = rgb_to_bgr(img).numpy()
        elif img.shape[0] == 4 and rgb2bgr:  # RGBA
            # RGBA to BGRA -> in tensor, if using OpenCV, else not needed. Only if image has colors.
            img_np = rgba_to_bgra(img).numpy()
        else:
            img_np = img.numpy()
        img_np = np.transpose(img_np, (1, 2, 0))  # CHW to HWC
    elif n_dim == 2:
        img_np = img.numpy()
    else:
        raise TypeError(
            f"Only support 4D, 3D and 2D tensor. But received with dimension: {n_dim:d}"
        )

    # if rgb2bgr:
    # img_np = img_np[[2, 1, 0], :, :] #RGB to BGR -> in numpy, if using OpenCV, else not needed. Only if image has colors.
    # TODO: Check: could denormalize in the begining in tensor form instead
    if denormalize:
        img_np = np_denorm(img_np)  # denormalize if needed
    if change_range:
        img_np = np.clip(
            data_range * img_np, 0, data_range
        ).round()  # np.clip to the data_range

    # has to be in range (0,255) before changing to np.uint8, else np.float32
    return img_np.astype(imtype)


def safe_cuda_cache_empty():
    """
    Empties the CUDA cache if CUDA is available. Hopefully without causing any errors.
    """
    try:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except:
        pass

nodes/impl/pytorch/utils.py
vladmandic/sd-extension-chainner
"""
Rembg code is modified from and copyright of Daniel Gatis,
and can be found here: https://github.com/danielgatis/rembg
"""

nodes/impl/rembg/__init__.py
vladmandic/sd-extension-chainner
from typing import List, Tuple

import numpy as np
import onnxruntime as ort
from cv2 import (
    BORDER_DEFAULT,
    COLOR_BGR2RGB,
    COLOR_BGRA2RGB,
    COLOR_RGBA2BGRA,
    MORPH_ELLIPSE,
    MORPH_OPEN,
    GaussianBlur,
    cvtColor,
    getStructuringElement,
    morphologyEx,
)
from PIL import Image
from PIL.Image import Image as PILImage
from scipy.ndimage import binary_erosion

from ...impl.image_utils import normalize
from ...utils.utils import get_h_w_c
from .pymatting.estimate_alpha_cf import estimate_alpha_cf
from .pymatting.estimate_foreground_ml import estimate_foreground_ml
from .pymatting.util import stack_images
from .session_factory import new_session

kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))


def alpha_matting_cutout(
    img: PILImage,
    mask: PILImage,
    foreground_threshold: int,
    background_threshold: int,
    erode_structure_size: int,
) -> PILImage:
    if img.mode == "RGBA" or img.mode == "CMYK":
        img = img.convert("RGB")

    npimg = np.asarray(img)
    npmask = np.asarray(mask)

    is_foreground = npmask > foreground_threshold
    is_background = npmask < background_threshold

    structure = None
    if erode_structure_size > 0:
        structure = np.ones(
            (erode_structure_size, erode_structure_size), dtype=np.uint8
        )

    is_foreground = binary_erosion(is_foreground, structure=structure)
    is_background = binary_erosion(is_background, structure=structure, border_value=1)

    trimap = np.full(npmask.shape, dtype=np.uint8, fill_value=128)
    trimap[is_foreground] = 255
    trimap[is_background] = 0

    img_normalized = npimg / 255.0
    trimap_normalized = trimap / 255.0

    alpha = estimate_alpha_cf(img_normalized, trimap_normalized)
    foreground = estimate_foreground_ml(img_normalized, alpha)
    cutout = stack_images(foreground, alpha)

    cutout = np.clip(cutout * 255, 0, 255).astype(np.uint8)  # type: ignore
    cutout = Image.fromarray(cutout)

    return cutout


def naive_cutout(img: PILImage, mask: PILImage) -> PILImage:
    empty = Image.new("RGBA", (img.size), 0)
    cutout = Image.composite(img, empty, mask)
    return cutout


def get_concat_v_multi(imgs: List[PILImage]) -> PILImage:
    pivot = imgs.pop(0)
    for im in imgs:
        pivot = get_concat_v(pivot, im)
    return pivot


def get_concat_v(img1: PILImage, img2: PILImage) -> PILImage:
    dst = Image.new("RGBA", (img1.width, img1.height + img2.height))
    dst.paste(img1, (0, 0))
    dst.paste(img2, (0, img1.height))
    return dst


def post_process(mask: np.ndarray) -> np.ndarray:
    """
    Post Process the mask for a smooth boundary by applying Morphological Operations
    Research based on paper: https://www.sciencedirect.com/science/article/pii/S2352914821000757
    args:
        mask: Binary Numpy Mask
    """
    mask = morphologyEx(mask, MORPH_OPEN, kernel)
    mask = GaussianBlur(mask, (5, 5), sigmaX=2, sigmaY=2, borderType=BORDER_DEFAULT)
    mask = np.where(mask < 127, 0, 255).astype(  # type: ignore
        np.uint8
    )  # convert again to binary
    return mask


def remove_bg(
    img: np.ndarray,
    ort_session: ort.InferenceSession,
    alpha_matting: bool = False,
    alpha_matting_foreground_threshold: int = 240,
    alpha_matting_background_threshold: int = 10,
    alpha_matting_erode_size: int = 10,
    post_process_mask: bool = False,
) -> Tuple[np.ndarray, np.ndarray]:
    # Flip channels to RGB mode and convert to PIL Image
    img = (img * 255).astype(np.uint8)
    _, _, c = get_h_w_c(img)
    if c == 3:
        img = cvtColor(img, COLOR_BGR2RGB)
    elif c == 4:
        img = cvtColor(img, COLOR_BGRA2RGB)
    pimg = Image.fromarray(img)

    session = new_session(ort_session)

    masks = session.predict(pimg)
    cutouts = []

    assert len(masks) > 0, "Model failed to generate masks"

    for mask in masks:
        if post_process_mask:
            mask = Image.fromarray(post_process(np.array(mask)))

        if alpha_matting:
            try:
                cutout = alpha_matting_cutout(
                    pimg,
                    mask,
                    alpha_matting_foreground_threshold,
                    alpha_matting_background_threshold,
                    alpha_matting_erode_size,
                )
            except ValueError:
                cutout = naive_cutout(pimg, mask)
        else:
            cutout = naive_cutout(pimg, mask)

        cutouts.append(cutout)

    cutout = pimg
    if len(cutouts) > 0:
        cutout = get_concat_v_multi(cutouts)

    return cvtColor(normalize(np.asarray(cutout)), COLOR_RGBA2BGRA), normalize(np.asarray(mask))  # type: ignore  pylint: disable=undefined-loop-variable

nodes/impl/rembg/bg.py
vladmandic/sd-extension-chainner
"""
Pymatting code is modified from and copyright of PyMatting,
and can be found here: https://github.com/pymatting/pymatting
"""

nodes/impl/rembg/pymatting/__init__.py
vladmandic/sd-extension-chainner
import numpy as np
import scipy.sparse
from numba import njit

from ....utils.utils import get_h_w_c


@njit(
    "void(f8[:, :, :], f8, i8, f8[:, :, :], i8[:], i8[:], b1[:, :])",
    cache=True,
    nogil=True,
)
def _cf_laplacian(image, epsilon, r, values, indices, indptr, is_known):
    h, w, d = image.shape
    assert d == 3
    size = 2 * r + 1
    window_area = size * size

    for yi in range(h):
        for xi in range(w):
            i = xi + yi * w
            k = i * (4 * r + 1) ** 2
            for yj in range(yi - 2 * r, yi + 2 * r + 1):
                for xj in range(xi - 2 * r, xi + 2 * r + 1):
                    j = xj + yj * w

                    if 0 <= xj < w and 0 <= yj < h:
                        indices[k] = j

                    k += 1

            indptr[i + 1] = k

    # Centered and normalized window colors
    c = np.zeros((2 * r + 1, 2 * r + 1, 3))

    # For each pixel of image
    for y in range(r, h - r):
        for x in range(r, w - r):
            if np.all(is_known[y - r : y + r + 1, x - r : x + r + 1]):
                continue

            # For each color channel
            for dc in range(3):
                # Calculate sum of color channel in window
                s = 0.0
                for dy in range(size):
                    for dx in range(size):
                        s += image[y + dy - r, x + dx - r, dc]

                # Calculate centered window color
                for dy in range(2 * r + 1):
                    for dx in range(2 * r + 1):
                        c[dy, dx, dc] = (
                            image[y + dy - r, x + dx - r, dc] - s / window_area
                        )

            # Calculate covariance matrix over color channels with epsilon regularization
            a00 = epsilon
            a01 = 0.0
            a02 = 0.0
            a11 = epsilon
            a12 = 0.0
            a22 = epsilon

            for dy in range(size):
                for dx in range(size):
                    a00 += c[dy, dx, 0] * c[dy, dx, 0]
                    a01 += c[dy, dx, 0] * c[dy, dx, 1]
                    a02 += c[dy, dx, 0] * c[dy, dx, 2]
                    a11 += c[dy, dx, 1] * c[dy, dx, 1]
                    a12 += c[dy, dx, 1] * c[dy, dx, 2]
                    a22 += c[dy, dx, 2] * c[dy, dx, 2]

            a00 /= window_area
            a01 /= window_area
            a02 /= window_area
            a11 /= window_area
            a12 /= window_area
            a22 /= window_area

            det = (
                a00 * a12 * a12
                + a01 * a01 * a22
                + a02 * a02 * a11
                - a00 * a11 * a22
                - 2 * a01 * a02 * a12
            )

            inv_det = 1.0 / det

            # Calculate inverse covariance matrix
            m00 = (a12 * a12 - a11 * a22) * inv_det
            m01 = (a01 * a22 - a02 * a12) * inv_det
            m02 = (a02 * a11 - a01 * a12) * inv_det
            m11 = (a02 * a02 - a00 * a22) * inv_det
            m12 = (a00 * a12 - a01 * a02) * inv_det
            m22 = (a01 * a01 - a00 * a11) * inv_det

            # For each pair ((xi, yi), (xj, yj)) in a (2 r + 1)x(2 r + 1) window
            for dyi in range(2 * r + 1):
                for dxi in range(2 * r + 1):
                    s = c[dyi, dxi, 0]
                    t = c[dyi, dxi, 1]
                    u = c[dyi, dxi, 2]

                    c0 = m00 * s + m01 * t + m02 * u
                    c1 = m01 * s + m11 * t + m12 * u
                    c2 = m02 * s + m12 * t + m22 * u

                    for dyj in range(2 * r + 1):
                        for dxj in range(2 * r + 1):
                            xi = x + dxi - r
                            yi = y + dyi - r
                            xj = x + dxj - r
                            yj = y + dyj - r

                            i = xi + yi * w
                            j = xj + yj * w

                            # Calculate contribution of pixel pair to L_ij
                            temp = (
                                c0 * c[dyj, dxj, 0]
                                + c1 * c[dyj, dxj, 1]
                                + c2 * c[dyj, dxj, 2]
                            )

                            value = (1.0 if (i == j) else 0.0) - (
                                1 + temp
                            ) / window_area

                            dx = xj - xi + 2 * r
                            dy = yj - yi + 2 * r

                            values[i, dy, dx] += value


def cf_laplacian(image, epsilon=1e-7, radius=1, is_known=None):
    """
    This function implements the alpha estimator for closed-form alpha matting as proposed by :cite:`levin2007closed`.
    Parameters
    ------------
    image: numpy.ndarray
       Image with shape :math:`h\\times w \\times 3`
    epsilon: float
       Regularization strength, defaults to :math:`10^{-7}`. Strong regularization improves convergence but results in smoother alpha mattes.
    radius: int
       Radius of local window size, defaults to :math:`1`, i.e. only adjacent pixels are considered.
       The size of the local window is given as :math:`(2 r + 1)^2`, where :math:`r` denotes         the radius. A larger radius might lead to violated color line constraints, but also
       favors further propagation of information within the image.
    is_known: numpy.ndarray
        Binary mask of pixels for which to compute the laplacian matrix.
        Laplacian entries for known pixels will have undefined values.
    Returns
    -------
    L: scipy.sparse.spmatrix
        Matting Laplacian
    """
    h, w, _ = get_h_w_c(image)
    n = h * w

    if is_known is None:
        is_known = np.zeros((h, w), dtype=np.bool_)

    is_known = is_known.reshape(h, w)

    # Data for matting laplacian in csr format
    indptr = np.zeros(n + 1, dtype=np.int64)
    indices = np.zeros(n * (4 * radius + 1) ** 2, dtype=np.int64)
    values = np.zeros((n, 4 * radius + 1, 4 * radius + 1), dtype=np.float64)

    _cf_laplacian(image, epsilon, radius, values, indices, indptr, is_known)

    L = scipy.sparse.csr_matrix((values.ravel(), indices, indptr), (n, n))

    return L

nodes/impl/rembg/pymatting/cf_laplacian.py
vladmandic/sd-extension-chainner
import numpy as np


def cg(
    A,
    b,
    x0=None,
    atol=0.0,
    rtol=1e-7,
    maxiter=10000,
    callback=None,
    M=None,
    reorthogonalize=False,
):
    """Solves a system of linear equations :math:`Ax=b` using conjugate gradient descent :cite:`hestenes1952methods`

    Parameters
    ----------
    A: scipy.sparse.csr_matrix
       Square matrix
    b: numpy.ndarray
       Vector describing the right-hand side of the system
    x0: numpy.ndarray
       Initialization, if `None` then :code:`x=np.zeros_like(b)`
    atol: float
       Absolute tolerance. The loop terminates if the :math:`||r||` is smaller than `atol`, where :math:`r` denotes the residual of the current iterate.
    rtol: float
       Relative tolerance. The loop terminates if :math:`{||r||}/{||b||}` is smaller than `rtol`, where :math:`r` denotes the residual of the current iterate.
    callback: function
       Function :code:`callback(A, x, b, norm_b, r, norm_r)` called after each iteration, defaults to `None`
    M: function or scipy.sparse.csr_matrix
       Function that applies the preconditioner to a vector. Alternatively, `M` can be a matrix describing the precondioner.
    reorthogonalize: boolean
        Wether to apply reorthogonalization of the residuals after each update, defaults to `False`


    Returns
    -------
    x: numpy.ndarray
        Solution of the system

    Example
    -------
    >>> from pymatting import *
    >>> import numpy as np
    >>> A = np.array([[3.0, 1.0], [1.0, 2.0]])
    >>> M = jacobi(A)
    >>> b = np.array([4.0, 3.0])
    >>> cg(A, b, M=M)
    array([1., 1.])
    """
    if M is None:

        def precondition(x):
            return x

    elif callable(M):
        precondition = M
    else:

        def precondition(x):
            return M.dot(x)

    x = np.zeros_like(b) if x0 is None else x0.copy()

    norm_b = np.linalg.norm(b)

    if callable(A):
        r = b - A(x)
    else:
        r = b - A.dot(x)

    norm_r = np.linalg.norm(r)

    if norm_r < atol or norm_r < rtol * norm_b:
        return x

    z = precondition(r)
    p = z.copy()
    rz = np.inner(r, z)

    for _ in range(maxiter):
        r_old = r.copy()

        if callable(A):
            Ap = A(p)
        else:
            Ap = A.dot(p)

        alpha = rz / np.inner(p, Ap)
        x += alpha * p
        r -= alpha * Ap

        norm_r = np.linalg.norm(r)

        if callback is not None:
            callback(A, x, b, norm_b, r, norm_r)

        if norm_r < atol or norm_r < rtol * norm_b:
            return x

        z = precondition(r)

        if reorthogonalize:
            beta = np.inner(r - r_old, z) / rz
            rz = np.inner(r, z)
        else:
            beta = 1.0 / rz
            rz = np.inner(r, z)
            beta *= rz

        p *= beta
        p += z

    raise ValueError(
        "Conjugate gradient descent did not converge within %d iterations" % maxiter
    )

nodes/impl/rembg/pymatting/cg.py
vladmandic/sd-extension-chainner
import numpy as np

from .cf_laplacian import cf_laplacian
from .cg import cg
from .ichol import ichol
from .util import sanity_check_image, trimap_split


def estimate_alpha_cf(  # pylint: disable=dangerous-default-value
    image, trimap, preconditioner=None, laplacian_kwargs=None, cg_kwargs=None
):
    """
    Estimate alpha from an input image and an input trimap using Closed-Form Alpha Matting as proposed by :cite:`levin2007closed`.

    Parameters
    ----------
    image: numpy.ndarray
        Image with shape :math:`h \\times  w \\times d` for which the alpha matte should be estimated
    trimap: numpy.ndarray
        Trimap with shape :math:`h \\times  w` of the image
    preconditioner: function or scipy.sparse.linalg.LinearOperator
        Function or sparse matrix that applies the preconditioner to a vector (default: ichol)
    laplacian_kwargs: dictionary
        Arguments passed to the :code:`cf_laplacian` function
    cg_kwargs: dictionary
        Arguments passed to the :code:`cg` solver
    is_known: numpy.ndarray
        Binary mask of pixels for which to compute the laplacian matrix.
        Providing this parameter might improve performance if few pixels are unknown.

    Returns
    -------
    alpha: numpy.ndarray
        Estimated alpha matte

    Example
    -------
    >>> from pymatting import *
    >>> image = load_image("data/lemur/lemur.png", "RGB")
    >>> trimap = load_image("data/lemur/lemur_trimap.png", "GRAY")
    >>> alpha = estimate_alpha_cf(
    ...     image,
    ...     trimap,
    ...     laplacian_kwargs={"epsilon": 1e-6},
    ...     cg_kwargs={"maxiter":2000})
    """
    if cg_kwargs is None:
        cg_kwargs = {}
    if laplacian_kwargs is None:
        laplacian_kwargs = {}
    if preconditioner is None:
        preconditioner = ichol

    sanity_check_image(image)

    is_fg, _, is_known, is_unknown = trimap_split(trimap)

    L = cf_laplacian(image, **laplacian_kwargs, is_known=is_known)

    # Split Laplacian matrix L into
    #
    #     [L_U   R ]
    #     [R^T   L_K]
    #
    # and then solve L_U x_U = -R is_fg_K for x where K (is_known) corresponds to
    # fixed pixels and U (is_unknown) corresponds to unknown pixels. For reference, see
    # Grady, Leo, et al. "Random walks for interactive alpha-matting." Proceedings of VIIP. Vol. 2005. 2005.

    L_U = L[is_unknown, :][:, is_unknown]

    R = L[is_unknown, :][:, is_known]

    m = is_fg[is_known]

    x = trimap.copy().ravel()

    x[is_unknown] = cg(L_U, -R.dot(m), M=preconditioner(L_U), **cg_kwargs)

    alpha = np.clip(x, 0, 1).reshape(trimap.shape)

    return alpha

nodes/impl/rembg/pymatting/estimate_alpha_cf.py
vladmandic/sd-extension-chainner
import numpy as np
from numba import njit, prange


@njit("void(f4[:, :, :], f4[:, :, :])", cache=True, nogil=True, parallel=True)
def _resize_nearest_multichannel(dst, src):
    """
    Internal method.

    Resize image src to dst using nearest neighbors filtering.
    Images must have multiple color channels, i.e. :code:`len(shape) == 3`.

    Parameters
    ----------
    dst: numpy.ndarray of type np.float32
        output image
    src: numpy.ndarray of type np.float32
        input image
    """
    h_src, w_src, depth = src.shape
    h_dst, w_dst, depth = dst.shape

    for y_dst in prange(h_dst):  # pylint: disable=not-an-iterable
        for x_dst in range(w_dst):
            x_src = max(0, min(w_src - 1, x_dst * w_src // w_dst))
            y_src = max(0, min(h_src - 1, y_dst * h_src // h_dst))

            for c in range(depth):
                dst[y_dst, x_dst, c] = src[y_src, x_src, c]


@njit("void(f4[:, :], f4[:, :])", cache=True, nogil=True, parallel=True)
def _resize_nearest(dst, src):
    """
    Internal method.

    Resize image src to dst using nearest neighbors filtering.
    Images must be grayscale, i.e. :code:`len(shape) == 3`.

    Parameters
    ----------
    dst: numpy.ndarray of type np.float32
        output image
    src: numpy.ndarray of type np.float32
        input image
    """
    h_src, w_src = src.shape
    h_dst, w_dst = dst.shape

    for y_dst in prange(h_dst):  # pylint: disable=not-an-iterable
        for x_dst in range(w_dst):
            x_src = max(0, min(w_src - 1, x_dst * w_src // w_dst))
            y_src = max(0, min(h_src - 1, y_dst * h_src // h_dst))

            dst[y_dst, x_dst] = src[y_src, x_src]


# TODO
# There should be an option to switch @njit(parallel=True) on or off.
# parallel=True would be faster, but might cause race conditions.
# User should have the option to turn it on or off.
@njit(
    "Tuple((f4[:, :, :], f4[:, :, :]))(f4[:, :, :], f4[:, :], f4, i4, i4, i4, f4)",
    cache=True,
    nogil=True,
)
def _estimate_fb_ml(
    input_image,
    input_alpha,
    regularization,
    n_small_iterations,
    n_big_iterations,
    small_size,
    gradient_weight,
):
    h0, w0, depth = input_image.shape

    dtype = np.float32

    w_prev = 1
    h_prev = 1

    F_prev = np.empty((h_prev, w_prev, depth), dtype=dtype)
    B_prev = np.empty((h_prev, w_prev, depth), dtype=dtype)

    n_levels = int(np.ceil(np.log2(max(w0, h0))))

    F = B = None
    for i_level in range(n_levels + 1):
        w = round(w0 ** (i_level / n_levels))
        h = round(h0 ** (i_level / n_levels))

        image = np.empty((h, w, depth), dtype=dtype)
        alpha = np.empty((h, w), dtype=dtype)

        _resize_nearest_multichannel(image, input_image)
        _resize_nearest(alpha, input_alpha)

        F = np.empty((h, w, depth), dtype=dtype)
        B = np.empty((h, w, depth), dtype=dtype)

        _resize_nearest_multichannel(F, F_prev)
        _resize_nearest_multichannel(B, B_prev)

        if w <= small_size and h <= small_size:
            n_iter = n_small_iterations
        else:
            n_iter = n_big_iterations

        b = np.zeros((2, depth), dtype=dtype)

        dx = [-1, 1, 0, 0]
        dy = [0, 0, -1, 1]

        for _ in range(n_iter):
            for y in prange(h):  # pylint: disable=not-an-iterable
                for x in range(w):
                    a0 = alpha[y, x]
                    a1 = 1.0 - a0

                    a00 = a0 * a0
                    a01 = a0 * a1
                    # a10 = a01 can be omitted due to symmetry of matrix
                    a11 = a1 * a1

                    for c in range(depth):
                        b[0, c] = a0 * image[y, x, c]
                        b[1, c] = a1 * image[y, x, c]

                    for d in range(4):
                        x2 = max(0, min(w - 1, x + dx[d]))
                        y2 = max(0, min(h - 1, y + dy[d]))

                        gradient = abs(a0 - alpha[y2, x2])

                        da = regularization + gradient_weight * gradient

                        a00 += da
                        a11 += da

                        for c in range(depth):
                            b[0, c] += da * F[y2, x2, c]
                            b[1, c] += da * B[y2, x2, c]

                    determinant = a00 * a11 - a01 * a01

                    inv_det = 1.0 / determinant

                    b00 = inv_det * a11
                    b01 = inv_det * -a01
                    b11 = inv_det * a00

                    for c in range(depth):
                        F_c = b00 * b[0, c] + b01 * b[1, c]
                        B_c = b01 * b[0, c] + b11 * b[1, c]

                        F_c = max(0.0, min(1.0, F_c))
                        B_c = max(0.0, min(1.0, B_c))

                        F[y, x, c] = F_c
                        B[y, x, c] = B_c

        F_prev = F
        B_prev = B

        w_prev = w
        h_prev = h

    if F is None or B is None:
        raise ValueError("Input image does not have multiple levels.")

    return F, B


def estimate_foreground_ml(
    image,
    alpha,
    regularization=1e-5,
    n_small_iterations=10,
    n_big_iterations=2,
    small_size=32,
    return_background=False,
    gradient_weight=1.0,
):
    """Estimates the foreground of an image given its alpha matte.

    See :cite:`germer2020multilevel` for reference.

    Parameters
    ----------
    image: numpy.ndarray
        Input image with shape :math:`h \\times  w \\times d`
    alpha: numpy.ndarray
        Input alpha matte shape :math:`h \\times  w`
    regularization: float
        Regularization strength :math:`\\epsilon`, defaults to :math:`10^{-5}`.
        Higher regularization results in smoother colors.
    n_small_iterations: int
        Number of iterations performed on small scale, defaults to :math:`10`
    n_big_iterations: int
        Number of iterations performed on large scale, defaults to :math:`2`
    small_size: int
        Threshold that determines at which size `n_small_iterations` should be used
    return_background: bool
        Whether to return the estimated background in addition to the foreground
    gradient_weight: float
        Larger values enforce smoother foregrounds, defaults to :math:`1`

    Returns
    -------
    F: numpy.ndarray
        Extracted foreground
    B: numpy.ndarray
        Extracted background

    Example
    -------
    >>> from pymatting import *
    >>> image = load_image("data/lemur/lemur.png", "RGB")
    >>> alpha = load_image("data/lemur/lemur_alpha.png", "GRAY")
    >>> F = estimate_foreground_ml(image, alpha, return_background=False)
    >>> F, B = estimate_foreground_ml(image, alpha, return_background=True)

    See Also
    ----
    stack_images: This function can be used to place the foreground on a new background.
    """

    foreground, background = _estimate_fb_ml(
        image.astype(np.float32),
        alpha.astype(np.float32),
        regularization,
        n_small_iterations,
        n_big_iterations,
        small_size,
        gradient_weight,
    )

    if return_background:
        return foreground, background

    return foreground

nodes/impl/rembg/pymatting/estimate_foreground_ml.py
vladmandic/sd-extension-chainner
import numpy as np
import scipy.sparse
from numba import njit


@njit(
    "i8(i8, f8[:], i8[:], i8[:], f8[:], i8[:], i8[:], f8, f8, i8, f8, b1)",
    cache=True,
    nogil=True,
)
def _ichol(
    n,
    Av,
    Ar,
    Ap,
    Lv,
    Lr,
    Lp,
    discard_threshold,
    shift,
    max_nnz,
    relative_discard_threshold,
    diag_keep_discarded,
):
    """
    :cite:`jones1995improved` might be slightly interesting for the general idea
    to use linked list to keep track of the sparse matrix values. But instead of
    pointers, we have to use indices here, since this is Python and not C.
    """
    nnz = 0
    c_n = 0
    s = np.zeros(n, np.int64)  # Next non-zero row index i in column j of L
    t = np.zeros(n, np.int64)  # First subdiagonal index i in column j of A
    l = (
        np.zeros(n, np.int64) - 1
    )  # Linked list of non-zero columns in row k of L; type: ignore
    a = np.zeros(n, np.float64)  # Values of column j
    r = np.zeros(n, np.float64)  # r[j] = sum(abs(A[j:, j])) for relative threshold
    b = np.zeros(
        n, np.bool_
    )  # b[i] indicates if the i-th element of column j is non-zero
    c = np.empty(n, np.int64)  # Row indices of non-zero elements in column j
    d = np.full(n, shift, np.float64)  # Diagonal elements of A
    for j in range(n):
        for idx in range(Ap[j], Ap[j + 1]):
            i = Ar[idx]
            if i == j:
                d[j] += Av[idx]
                t[j] = idx + 1
            if i >= j:
                r[j] += abs(Av[idx])
    for j in range(n):  # For each column j
        for idx in range(t[j], Ap[j + 1]):  # For each L_ij
            i = Ar[idx]
            L_ij = Av[idx]
            if L_ij != 0.0 and i > j:
                a[i] += L_ij  # Assign non-zero value to L_ij in sparse column
                if not b[i]:
                    b[i] = True  # Mark it as non-zero
                    c[c_n] = i  # Remember index for later deletion
                    c_n += 1
        k = l[j]  # Find index k of column with non-zero element in row j
        while k != -1:  # For each column of that type
            k0 = s[k]  # Start index of non-zero elements in column k
            k1 = Lp[k + 1]  # End index
            k2 = l[k]  # Remember next column index before it is overwritten
            L_jk = Lv[k0]  # Value of non-zero element at start of column
            k0 += 1  # Advance to next non-zero element in column
            if k0 < k1:  # If there is a next non-zero element
                s[k] = k0  # Advance start index in column k to next non-zero element
                i = Lr[k0]  # Row index of next non-zero element in column k
                l[k] = l[i]  # Remember old list i index in list k
                l[i] = k  # Insert index of non-zero element into list i
                for idx in range(k0, k1):  # For each non-zero L_ik in column k
                    i = Lr[idx]
                    L_ik = Lv[idx]
                    a[i] -= L_ik * L_jk  # Update element L_ij in sparse column
                    if not b[i]:  # Check if sparse column element was zero
                        b[i] = True  # Mark as non-zero in sparse column
                        c[c_n] = i  # Remember index for later deletion
                        c_n += 1
            k = k2  # Advance to next column k
        if d[j] <= 0.0:
            return -1
        if nnz + 1 + c_n > max_nnz:
            return -2
        d[j] = np.sqrt(d[j])  # Update diagonal element L_ii
        Lv[nnz] = d[j]  # Add diagonal element L_ii to L
        Lr[nnz] = j  # Add row index of L_ii to L
        nnz += 1
        s[j] = nnz  # Set first non-zero index of column j
        for i in np.sort(
            c[:c_n]
        ):  # Sort row indices of column j for correct insertion order into L
            L_ij = a[i] / d[j]  # Get non-zero element from sparse column j
            if diag_keep_discarded:
                d[i] -= L_ij * L_ij  # Update diagonal element L_ii
            rel = (
                relative_discard_threshold * r[j]
            )  # Relative discard threshold (before div)
            if (
                abs(L_ij) > discard_threshold and abs(a[i]) > rel
            ):  # If element is sufficiently non-zero
                if not diag_keep_discarded:
                    d[i] -= L_ij * L_ij  # Update diagonal element L_ii
                Lv[nnz] = L_ij  # Add element L_ij to L
                Lr[nnz] = i  # Add row index of L_ij
                nnz += 1
            a[i] = 0.0  # Set element i in column j to zero
            b[i] = False  # Mark element as zero
        c_n = 0  # Discard row indices of non-zero elements in column j.
        Lp[j + 1] = nnz  # Update count of non-zero elements up to column j
        if Lp[j] + 1 < Lp[j + 1]:  # If column j has a non-zero element below diagonal
            i = Lr[Lp[j] + 1]  # Row index of first off-diagonal non-zero element
            l[j] = l[i]  # Remember old list i index in list j
            l[i] = j  # Insert index of non-zero element into list i
    return nnz


@njit("void(f8[:], i8[:], i8[:], f8[:], i8)", cache=True, nogil=True)
def _backsub_L_csc_inplace(data, indices, indptr, x, n):
    for j in range(n):
        k = indptr[j]
        L_jj = data[k]
        temp = x[j] / L_jj

        x[j] = temp

        for k in range(indptr[j] + 1, indptr[j + 1]):
            i = indices[k]
            L_ij = data[k]

            x[i] -= L_ij * temp


@njit("void(f8[:], i8[:], i8[:], f8[:], i8)", cache=True, nogil=True)
def _backsub_LT_csc_inplace(data, indices, indptr, x, n):
    for i in range(n - 1, -1, -1):
        s = x[i]

        for k in range(indptr[i] + 1, indptr[i + 1]):
            j = indices[k]
            L_ji = data[k]
            s -= L_ji * x[j]

        k = indptr[i]
        L_ii = data[k]

        x[i] = s / L_ii


class CholeskyDecomposition(object):
    """Cholesky Decomposition

    Calling this object applies the preconditioner to a vector by forward and back substitution.

    Parameters
    ----------
    Ltuple: tuple of numpy.ndarrays
        Tuple of array describing values, row indices and row pointers for Cholesky factor in the compressed sparse column format (csc)
    """

    def __init__(self, Ltuple):
        self.Ltuple = Ltuple

    @property
    def L(self):
        """Returns the Cholesky factor

        Returns
        -------
        L: scipy.sparse.csc_matrix
            Cholesky factor
        """
        _, _, Lp = self.Ltuple
        n = len(Lp) - 1
        return scipy.sparse.csc_matrix(self.Ltuple, (n, n))

    def __call__(self, b):
        Lv, Lr, Lp = self.Ltuple
        n = len(b)
        x = b.copy()
        _backsub_L_csc_inplace(Lv, Lr, Lp, x, n)
        _backsub_LT_csc_inplace(Lv, Lr, Lp, x, n)
        return x


def ichol(  # pylint: disable=dangerous-default-value
    A,
    discard_threshold=1e-4,
    shifts=None,
    max_nnz=int(4e9 / 16),
    relative_discard_threshold=0.0,
    diag_keep_discarded=True,
):
    """Implements the thresholded incomplete Cholesky decomposition

    Parameters
    ----------
    A: scipy.sparse.csc_matrix
        Matrix for which the preconditioner should be calculated
    discard_threshold: float
        Values having an absolute value smaller than this threshold will be discarded while calculating the Cholesky decompositions
    shifts: array of floats
        Values to try for regularizing the matrix of interest in case it is not positive definite after discarding the small values
    max_nnz: int
        Maximum number of non-zero entries in the Cholesky decomposition. Defaults to 250 million, which should usually be around 4 GB.
    relative_discard_threshold: float
        Values with an absolute value of less than :code:`relative_discard_threshold * sum(abs(A[j:, j]))` will be discarded.
        A dense ichol implementation with relative threshold would look like this::

            L = np.tril(A)
            for j in range(n):
                col = L[j:, j]
                col -= np.sum(L[j, :j] * L[j:, :j], axis=1)
                discard_mask = abs(col[1:]) < relative_discard_threshold * np.sum(np.abs(A[j:, j]))
                col[1:][discard_mask] = 0
                col[0] **= 0.5
                col[1:] /= col[0]

    diag_keep_discarded: bool
        Whether to update the diagonal with the discarded values. Usually better if :code:`True`.

    Returns
    -------
    chol: CholeskyDecomposition
        Preconditioner or solver object.

    Raises
    ------
    ValueError
        If inappropriate parameter values were passed

    Example
    -------
    >>> from pymatting import *
    >>> import numpy as np
    >>> from scipy.sparse import csc_matrix
    >>> A = np.array([[2.0, 3.0], [3.0, 5.0]])
    >>> cholesky_decomposition = ichol(csc_matrix(A))
    >>> cholesky_decomposition(np.array([1.0, 2.0]))
    array([-1.,  1.])
    """

    if shifts is None:
        shifts = [0.0, 0.0001, 0.001, 0.01, 0.1, 0.5, 1.0, 10.0, 100, 1000.0, 10000.0, 100000.0]
    if isinstance(A, scipy.sparse.csr_matrix):
        A = A.T

    if not isinstance(A, scipy.sparse.csc_matrix):
        raise ValueError("Matrix A must be a scipy.sparse.csc_matrix")

    if not A.has_canonical_format:
        A.sum_duplicates()

    m, n = A.shape

    assert m == n

    Lv = np.empty(max_nnz, dtype=np.float64)  # Values of non-zero elements of L
    Lr = np.empty(max_nnz, dtype=np.int64)  # Row indices of non-zero elements of L
    Lp = np.zeros(
        n + 1, dtype=np.int64
    )  # Start(Lp[i]) and end(Lp[i+1]) index of L[:, i] in Lv

    nnz = -3
    for shift in shifts:
        nnz = _ichol(
            n,
            A.data,
            A.indices.astype(np.int64),
            A.indptr.astype(np.int64),
            Lv,
            Lr,
            Lp,
            discard_threshold,
            shift,
            max_nnz,
            relative_discard_threshold,
            diag_keep_discarded,
        )

        if nnz >= 0:
            break

        if nnz == -1:
            print("PERFORMANCE WARNING:")
            print(
                "Thresholded incomplete Cholesky decomposition failed due to insufficient positive-definiteness of matrix A with parameters:"
            )
            print("    discard_threshold = %e" % discard_threshold)
            print("    shift = %e" % shift)
            print("Try decreasing discard_threshold or start with a larger shift")
            print("")

        if nnz == -2:
            raise ValueError(
                "Thresholded incomplete Cholesky decomposition failed because more than max_nnz non-zero elements were created. Try increasing max_nnz or discard_threshold."
            )

    if nnz < 0:
        raise ValueError(
            "Thresholded incomplete Cholesky decomposition failed due to insufficient positive-definiteness of matrix A and diagonal shifts did not help."
        )

    Lv = Lv[:nnz]
    Lr = Lr[:nnz]

    return CholeskyDecomposition((Lv, Lr, Lp))

nodes/impl/rembg/pymatting/ichol.py
vladmandic/sd-extension-chainner
import warnings

import numpy as np


def sanity_check_image(image):
    """Performs a sanity check for input images. Image values should be in the
    range [0, 1], the `dtype` should be `np.float32` or `np.float64` and the
    image shape should be `(?, ?, 3)`.
    Parameters
    ----------
    image: numpy.ndarray
        Image with shape :math:`h \\times w \\times 3`
    Example
    -------
    >>> import numpy as np
    >>> from pymatting import check_image
    >>> image = (np.random.randn(64, 64, 2) * 255).astype(np.int32)
    >>> sanity_check_image(image)
    __main__:1: UserWarning: Expected RGB image of shape (?, ?, 3), but image.shape is (64, 64, 2).
    __main__:1: UserWarning: Image values should be in [0, 1], but image.min() is -933.
    __main__:1: UserWarning: Image values should be in [0, 1], but image.max() is 999.
    __main__:1: UserWarning: Unexpected image.dtype int32. Are you sure that you do not want to use np.float32 or np.float64 instead?
    """

    if len(image.shape) != 3 or image.shape[2] != 3:
        warnings.warn(
            "Expected RGB image of shape (?, ?, 3), but image.shape is %s."
            % str(image.shape),
            stacklevel=3,
        )

    min_value = image.min()
    max_value = image.max()

    if min_value < 0.0:
        warnings.warn(
            "Image values should be in [0, 1], but image.min() is %s." % min_value,
            stacklevel=3,
        )

    if max_value > 1.0:
        warnings.warn(
            "Image values should be in [0, 1], but image.max() is %s." % max_value,
            stacklevel=3,
        )

    if image.dtype not in [np.float32, np.float64]:
        warnings.warn(
            "Unexpected image.dtype %s. Are you sure that you do not want to use np.float32 or np.float64 instead?"
            % image.dtype,
            stacklevel=3,
        )


def stack_images(*images):
    """This function stacks images along the third axis.
    This is useful for combining e.g. rgb color channels or color and alpha channels.
    Parameters
    ----------
    *images: numpy.ndarray
        Images to be stacked.
    Returns
    -------
    image: numpy.ndarray
        Stacked images as numpy.ndarray
    Example
    -------
    >>> from pymatting.util.util import stack_images
    >>> import numpy as np
    >>> I = stack_images(np.random.rand(4,5,3), np.random.rand(4,5,3))
    >>> I.shape
    (4, 5, 6)
    """
    images = [
        (image if len(image.shape) == 3 else image[:, :, np.newaxis])
        for image in images
    ]
    return np.concatenate(images, axis=2)


def trimap_split(trimap, flatten=True, bg_threshold=0.1, fg_threshold=0.9):
    """This function splits the trimap into foreground pixels, background pixels, and unknown pixels.
    Foreground pixels are pixels where the trimap has values larger than or equal to `fg_threshold` (default: 0.9).
    Background pixels are pixels where the trimap has values smaller than or equal to `bg_threshold` (default: 0.1).
    Pixels with other values are assumed to be unknown.
    Parameters
    ----------
    trimap: numpy.ndarray
        Trimap with shape :math:`h \\times w`
    flatten: bool
        If true np.flatten is called on the trimap
    Returns
    -------
    is_fg: numpy.ndarray
        Boolean array indicating which pixel belongs to the foreground
    is_bg: numpy.ndarray
        Boolean array indicating which pixel belongs to the background
    is_known: numpy.ndarray
        Boolean array indicating which pixel is known
    is_unknown: numpy.ndarray
        Boolean array indicating which pixel is unknown
    bg_threshold: float
        Pixels with smaller trimap values will be considered background.
    fg_threshold: float
        Pixels with larger trimap values will be considered foreground.
    Example
    -------
    >>> import numpy as np
    >>> from pymatting import *
    >>> trimap = np.array([[1,0],[0.5,0.2]])
    >>> is_fg, is_bg, is_known, is_unknown = trimap_split(trimap)
    >>> is_fg
    array([ True, False, False, False])
    >>> is_bg
    array([False,  True, False, False])
    >>> is_known
    array([ True,  True, False, False])
    >>> is_unknown
    array([False, False,  True,  True])
    """
    if flatten:
        trimap = trimap.flatten()

    min_value = trimap.min()
    max_value = trimap.max()

    if min_value < 0.0:
        warnings.warn(
            "Trimap values should be in [0, 1], but trimap.min() is %s." % min_value,
            stacklevel=3,
        )

    if max_value > 1.0:
        warnings.warn(
            "Trimap values should be in [0, 1], but trimap.max() is %s." % min_value,
            stacklevel=3,
        )

    if trimap.dtype not in [np.float32, np.float64]:
        warnings.warn(
            "Unexpected trimap.dtype %s. Are you sure that you do not want to use np.float32 or np.float64 instead?"
            % trimap.dtype,
            stacklevel=3,
        )

    is_fg = trimap >= fg_threshold
    is_bg = trimap <= bg_threshold

    if is_bg.sum() == 0:
        raise ValueError(
            "Trimap did not contain background values (values <= %f)" % bg_threshold
        )

    if is_fg.sum() == 0:
        raise ValueError(
            "Trimap did not contain foreground values (values >= %f)" % fg_threshold
        )

    is_known = is_fg | is_bg
    is_unknown = ~is_known

    return is_fg, is_bg, is_known, is_unknown

nodes/impl/rembg/pymatting/util.py
vladmandic/sd-extension-chainner
from typing import Dict, List, Tuple

import numpy as np
import onnxruntime as ort
from PIL import Image
from PIL.Image import Image as PILImage


class BaseSession:
    def __init__(
        self,
        inner_session: ort.InferenceSession,
        mean: Tuple[float, float, float],
        std: Tuple[float, float, float],
        size: Tuple[int, int],
    ):
        self.inner_session = inner_session
        self.mean = mean
        self.std = std
        self.size = size

    def normalize(self, img: PILImage) -> Dict[str, np.ndarray]:
        im = img.convert("RGB").resize(self.size, Image.LANCZOS)
        im_ary = np.array(im)
        im_ary = im_ary / np.max(im_ary)

        tmpImg = np.zeros((im_ary.shape[0], im_ary.shape[1], 3))
        tmpImg[:, :, 0] = (im_ary[:, :, 0] - self.mean[0]) / self.std[0]
        tmpImg[:, :, 1] = (im_ary[:, :, 1] - self.mean[1]) / self.std[1]
        tmpImg[:, :, 2] = (im_ary[:, :, 2] - self.mean[2]) / self.std[2]

        tmpImg = tmpImg.transpose((2, 0, 1))

        model_input_name = self.inner_session.get_inputs()[0].name

        return {model_input_name: np.expand_dims(tmpImg, 0).astype(np.float32)}

    def predict(self, _: PILImage) -> List[PILImage]:
        raise NotImplementedError

nodes/impl/rembg/session_base.py
vladmandic/sd-extension-chainner
from typing import List

import numpy as np
from PIL import Image
from PIL.Image import Image as PILImage
from scipy.special import log_softmax

from .session_base import BaseSession

pallete1 = [
    0,
    0,
    0,
    255,
    255,
    255,
    0,
    0,
    0,
    0,
    0,
    0,
]

pallete2 = [
    0,
    0,
    0,
    0,
    0,
    0,
    255,
    255,
    255,
    0,
    0,
    0,
]

pallete3 = [
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    255,
    255,
    255,
]


class ClothSession(BaseSession):
    def predict(self, img: PILImage) -> List[PILImage]:
        ort_outs = self.inner_session.run(None, self.normalize(img))

        pred = ort_outs
        pred = log_softmax(pred[0], 1)
        pred = np.argmax(pred, axis=1, keepdims=True)
        pred = np.squeeze(pred, 0)
        pred = np.squeeze(pred, 0)

        mask = Image.fromarray(pred.astype("uint8"), mode="L")
        mask = mask.resize(img.size, Image.LANCZOS)

        masks = []

        mask1 = mask.copy()
        mask1.putpalette(pallete1)
        mask1 = mask1.convert("RGB").convert("L")
        masks.append(mask1)

        mask2 = mask.copy()
        mask2.putpalette(pallete2)
        mask2 = mask2.convert("RGB").convert("L")
        masks.append(mask2)

        mask3 = mask.copy()
        mask3.putpalette(pallete3)
        mask3 = mask3.convert("RGB").convert("L")
        masks.append(mask3)

        return masks

nodes/impl/rembg/session_cloth.py
vladmandic/sd-extension-chainner
from typing import Type

import onnxruntime as ort

from ...impl.onnx.utils import get_input_shape
from .session_base import BaseSession
from .session_cloth import ClothSession
from .session_simple import SimpleSession


def new_session(session: ort.InferenceSession) -> BaseSession:
    session_class: Type[BaseSession]

    input_width = get_input_shape(session)[2]

    # Using size to determine session type and norm parameters is fragile,
    # but at the moment I don't know a better way to detect architecture due
    # to the lack of consistency in naming and outputs across arches and repos.
    # It works right now because of the limited number of models supported,
    # but if that expands, it may become necessary to find an alternative.
    mean = (0.485, 0.456, 0.406)
    std = (0.229, 0.224, 0.225)
    size = (input_width, input_width) if input_width is not None else (320, 320)
    if input_width == 768:  # U2NET cloth model
        session_class = ClothSession
        mean = (0.5, 0.5, 0.5)
        std = (0.5, 0.5, 0.5)
    else:
        session_class = SimpleSession
        if input_width == 1024:  # ISNET
            mean = (0.5, 0.5, 0.5)
            std = (1, 1, 1)
        elif input_width == 512:  # Models trained using anime-segmentation repo
            mean = (0, 0, 0)
            std = (1, 1, 1)

    return session_class(session, mean, std, size)

nodes/impl/rembg/session_factory.py
vladmandic/sd-extension-chainner
from typing import List

import numpy as np
from PIL import Image
from PIL.Image import Image as PILImage

from .session_base import BaseSession


class SimpleSession(BaseSession):
    def predict(self, img: PILImage) -> List[PILImage]:
        ort_outs = self.inner_session.run(None, self.normalize(img))

        pred = ort_outs[0][:, 0, :, :]

        ma = np.max(pred)
        mi = np.min(pred)

        pred = (pred - mi) / (ma - mi)
        pred = np.squeeze(pred)

        mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
        mask = mask.resize(img.size, Image.LANCZOS)

        return [mask]

nodes/impl/rembg/session_simple.py
vladmandic/sd-extension-chainner
import math
from enum import Enum

import cv2
import numpy as np

from ..utils.utils import get_h_w_c


class TileMode(Enum):
    TILE = 0
    MIRROR = 1


def tile_image(img: np.ndarray, width: int, height: int, mode: TileMode) -> np.ndarray:
    if mode == TileMode.TILE:
        # do nothing
        pass
    elif mode == TileMode.MIRROR:
        # flip the image to create a mirrored tile
        flip_x: np.ndarray = cv2.flip(img, 0)
        flip_y: np.ndarray = cv2.flip(img, 1)
        flip_xy: np.ndarray = cv2.flip(img, -1)

        img = cv2.vconcat(
            [
                cv2.hconcat([img, flip_y]),  # type: ignore
                cv2.hconcat([flip_x, flip_xy]),  # type: ignore
            ]
        )
    else:
        raise AssertionError(f"Invalid tile mode {mode}")

    h, w, _ = get_h_w_c(img)
    tile_w = math.ceil(width / w)
    tile_h = math.ceil(height / h)
    img = np.tile(img, (tile_h, tile_w) if img.ndim == 2 else (tile_h, tile_w, 1))

    # crop to make sure the dimensions are correct
    return img[:height, :width]

nodes/impl/tile.py
vladmandic/sd-extension-chainner
from __future__ import annotations
import math
from typing import Callable, Optional, Union
import numpy as np
from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
from nodes.log import logger, console
from ...utils.utils import Region, Size, get_h_w_c
from .exact_split import exact_split
from .tiler import Tiler


class Split:
    pass


SplitImageOp = Callable[[np.ndarray, Region], Union[np.ndarray, Split]]


def auto_split(
    img: np.ndarray,
    upscale: SplitImageOp,
    tiler: Tiler,
    overlap: int = 16,
) -> np.ndarray:
    """
    Splits the image into tiles according to the given tiler.

    This method only changes the size of the given image, the tiles passed into the upscale function will have same number of channels.

    The region passed into the upscale function is the region of the current tile.
    The size of the region is guaranteed to be the same as the size of the given tile.

    ## Padding

    If the given tiler allows smaller tile sizes, then it is guaranteed that no padding will be added.
    Otherwise, no padding is only guaranteed if the starting tile size is not larger than the size of the given image.
    """

    h, w, c = get_h_w_c(img)
    split = _max_split if tiler.allow_smaller_tile_size() else _exact_split

    return split(
        img,
        upscale=upscale,
        starting_tile_size=tiler.starting_tile_size(w, h, c),
        split_tile_size=tiler.split,
        overlap=overlap,
    )


class _SplitEx(Exception):
    pass


def _exact_split(
    img: np.ndarray,
    upscale: SplitImageOp,
    starting_tile_size: Size,
    split_tile_size: Callable[[Size], Size],
    overlap: int,
) -> np.ndarray:
    h, w, c = get_h_w_c(img)
    logger.info(f"chaiNNer: exact size split image ({w}x{h}px @ {c}) with exact tile size {starting_tile_size[0]}x{starting_tile_size[1]}px.")

    def no_split_upscale(i: np.ndarray, r: Region) -> np.ndarray:
        result = upscale(i, r)
        if isinstance(result, Split):
            raise _SplitEx
        return result

    MAX_ITER = 20
    for _ in range(MAX_ITER):
        try:
            max_overlap = min(*starting_tile_size) // 4
            return exact_split(
                img=img,
                exact_size=starting_tile_size,
                upscale=no_split_upscale,
                overlap=min(max_overlap, overlap),
            )
        except _SplitEx:
            starting_tile_size = split_tile_size(starting_tile_size)

    raise ValueError(f"Aborting after {MAX_ITER} splits. Unable to upscale image.")


def _max_split(
    img: np.ndarray,
    upscale: SplitImageOp,
    starting_tile_size: Size,
    split_tile_size: Callable[[Size], Size],
    overlap: int,
) -> np.ndarray:
    """
    Splits the image into tiles with at most the given tile size.

    If the upscale method requests a split, then the tile size will be lowered.
    """

    h, w, c = get_h_w_c(img)

    img_region = Region(0, 0, w, h)

    max_tile_size = starting_tile_size
    # logger.debug(f"chaiNNer: auto split image ({w}x{h}px @ {c}) with initial tile size {max_tile_size}")

    if w <= max_tile_size[0] and h <= max_tile_size[1]:
        # the image might be small enough so that we don't have to split at all
        upscale_result = upscale(img, img_region)
        if not isinstance(upscale_result, Split):
            return upscale_result
        # the image was too large
        max_tile_size = split_tile_size(max_tile_size)
        logger.info(f"chaiNNer: unable to upscale the whole image at once. Reduced tile size to {max_tile_size}")

    # The upscale method is allowed to request splits at any time.
    # When a split occurs, we have to "restart" the loop and
    # these 2 variables allow us to split the already processed tiles.
    start_x = 0
    start_y = 0

    # To allocate the result image, we need to know the upscale factor first,
    # and we only get to know this factor after the first successful upscale.
    result: Optional[np.ndarray] = None
    scale: int = 0

    restart = True
    while restart:
        restart = False

        # This is a bit complex.
        # We don't actually use the current tile size to partition the image.
        # If we did, then tile_size=1024 and w=1200 would result in very uneven tiles.
        # Instead, we use tile_size to calculate how many tiles we get in the x and y direction
        # and then calculate the optimal tile size for the x and y direction using the counts.
        # This yields optimal tile sizes which should prevent unnecessary splitting.
        tile_count_x = math.ceil(w / max_tile_size[0])
        tile_count_y = math.ceil(h / max_tile_size[1])
        tile_size_x = math.ceil(w / tile_count_x)
        tile_size_y = math.ceil(h / tile_count_y)
        # logger.debug(f"chaiNNer: Currently {tile_count_x}x{tile_count_y} tiles each {tile_size_x}x{tile_size_y}px.")
        with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=console) as progress:
            task = progress.add_task(description="Upscaling", total=tile_count_y * tile_count_x)
            for y in range(0, tile_count_y):
                if restart:
                    break
                if y < start_y:
                    continue
                for x in range(0, tile_count_x):
                    if y == start_y and x < start_x:
                        continue
                    tile = Region(x * tile_size_x, y * tile_size_y, tile_size_x, tile_size_y).intersect(img_region)
                    pad = img_region.child_padding(tile).min(overlap)
                    padded_tile = tile.add_padding(pad)
                    upscale_result = upscale(padded_tile.read_from(img), padded_tile)
                    if isinstance(upscale_result, Split):
                        max_tile_size = split_tile_size(max_tile_size)
                        new_tile_count_x = math.ceil(w / max_tile_size[0])
                        new_tile_count_y = math.ceil(h / max_tile_size[1])
                        new_tile_size_x = math.ceil(w / new_tile_count_x)
                        new_tile_size_y = math.ceil(h / new_tile_count_y)
                        start_x = (x * tile_size_x) // new_tile_size_x
                        start_y = (y * tile_size_x) // new_tile_size_y
                        logger.debug(f"chaiNNer: Split occurred. New tile size is {max_tile_size}. Starting at {start_x},{start_y}.")
                        restart = True
                        break
                    # figure out by how much the image was upscaled by
                    up_h, up_w, _ = get_h_w_c(upscale_result)
                    current_scale = up_h // padded_tile.height
                    assert current_scale > 0
                    assert padded_tile.height * current_scale == up_h
                    assert padded_tile.width * current_scale == up_w
                    if result is None:
                        # allocate the result image
                        scale = current_scale
                        result = np.zeros((h * scale, w * scale, c), dtype=np.float32)
                    assert current_scale == scale
                    # remove overlap padding
                    upscale_result = pad.scale(scale).remove_from(upscale_result)
                    # copy into result image
                    tile.scale(scale).write_into(result, upscale_result)
                    progress.update(task, advance=1, description="Upscaling")

    assert result is not None
    return result

nodes/impl/upscale/auto_split.py
vladmandic/sd-extension-chainner
from typing import Callable, NewType

import numpy as np
from nodes.log import logger

from ...utils.utils import get_h_w_c
from .tiler import MaxTileSize, NoTiling, Tiler


def estimate_tile_size(
    budget: int,
    model_size: int,
    img: np.ndarray,
    img_element_size: int = 4,
) -> int:
    h, w, c = get_h_w_c(img)
    img_bytes = h * w * c * img_element_size
    mem_required_estimation = (model_size / (1024 * 52)) * img_bytes

    tile_pixels = w * h * budget / mem_required_estimation
    # the largest power-of-2 tile_size such that tile_size**2 < tile_pixels
    tile_size = 2 ** (int(tile_pixels**0.5).bit_length() - 1)

    GB_AMT = 1024**3
    required_mem = f"{mem_required_estimation/GB_AMT:.2f}"
    budget_mem = f"{budget/GB_AMT:.2f}"
    logger.info(f"chaiNNer: estimating memory required: {required_mem} GB, {budget_mem} GB free Estimated tile size: {tile_size}")
    return tile_size


TileSize = NewType("TileSize", int)
ESTIMATE = TileSize(0)
NO_TILING = TileSize(-1)
MAX_TILE_SIZE = TileSize(-2)
TILE_SIZE_256 = TileSize(256)


def parse_tile_size_input(tile_size: TileSize, estimate: Callable[[], Tiler]) -> Tiler:
    if tile_size == 0:
        return estimate()
    if tile_size == -1:
        return NoTiling()
    if tile_size == -2:
        return MaxTileSize()

    assert tile_size > 0
    return MaxTileSize(tile_size)

nodes/impl/upscale/auto_split_tiles.py
vladmandic/sd-extension-chainner
from typing import Tuple

import numpy as np

from ...utils.utils import get_h_w_c
from ..image_op import ImageOp, clipped
from ..image_utils import as_target_channels


def with_black_and_white_backgrounds(img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    c = get_h_w_c(img)[2]
    assert c == 4

    black = np.copy(img[:, :, :3])
    white = np.copy(img[:, :, :3])
    for c in range(3):
        black[:, :, c] *= img[:, :, 3]
        white[:, :, c] = (white[:, :, c] - 1) * img[:, :, 3] + 1

    return black, white


def denoise_and_flatten_alpha(img: np.ndarray) -> np.ndarray:
    alpha_min = np.min(img, axis=2)
    alpha_max = np.max(img, axis=2)
    alpha_mean = np.mean(img, axis=2)
    alpha = alpha_max * alpha_mean + alpha_min * (1 - alpha_mean)
    return alpha


def convenient_upscale(
    img: np.ndarray,
    model_in_nc: int,
    model_out_nc: int,
    upscale: ImageOp,
    separate_alpha: bool = False,
) -> np.ndarray:
    """
    Upscales the given image in an intuitive/convenient way.

    This method guarantees that the `upscale` function will be called with an image with
    `model_in_nc` number of channels.

    Additionally, guarantees that the number of channels of the output image will match
    that of the input image in cases where `model_in_nc` == `model_out_nc`, and match
    `model_out_nc` otherwise.
    """
    in_img_c = get_h_w_c(img)[2]

    upscale = clipped(upscale)

    if model_in_nc != model_out_nc:
        return upscale(as_target_channels(img, model_in_nc, True))

    if in_img_c == model_in_nc:
        return upscale(img)

    if in_img_c == 4:
        # Ignore alpha if single-color or not being replaced
        unique = np.unique(img[:, :, 3])
        if len(unique) == 1:
            rgb = as_target_channels(
                upscale(as_target_channels(img[:, :, :3], model_in_nc, True)), 3, True
            )
            unique_alpha = np.full(rgb.shape[:-1], unique[0], np.float32)
            return np.dstack((rgb, unique_alpha))

        if separate_alpha:
            # Upscale the RGB channels and alpha channel separately
            rgb = as_target_channels(
                upscale(as_target_channels(img[:, :, :3], model_in_nc, True)), 3, True
            )
            alpha = denoise_and_flatten_alpha(
                upscale(as_target_channels(img[:, :, 3], model_in_nc, True))
            )
            return np.dstack((rgb, alpha))
        else:
            # Transparency hack (white/black background difference alpha)
            black, white = with_black_and_white_backgrounds(img)
            black_up = as_target_channels(
                upscale(as_target_channels(black, model_in_nc, True)), 3, True
            )
            white_up = as_target_channels(
                upscale(as_target_channels(white, model_in_nc, True)), 3, True
            )

            # Interpolate between the alpha values to get a more defined alpha
            alpha_candidates = 1 - (white_up - black_up)  #  type: ignore
            alpha = denoise_and_flatten_alpha(alpha_candidates)

            return np.dstack((black_up, alpha))

    return as_target_channels(
        upscale(as_target_channels(img, model_in_nc, True)), in_img_c, True
    )

nodes/impl/upscale/convenient_upscale.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Callable, List, Optional, Tuple

import numpy as np
from nodes.log import logger

from ...utils.utils import Padding, Region, Size, get_h_w_c
from ..image_utils import BorderType, create_border


def _pad_image(img: np.ndarray, min_size: Size):
    h, w, _ = get_h_w_c(img)

    min_w, min_h = min_size
    x = max(0, min_w - w) / 2
    y = max(0, min_h - h) / 2

    padding = Padding(math.floor(y), math.floor(x), math.ceil(y), math.ceil(x))

    return create_border(img, BorderType.REFLECT_MIRROR, padding), padding


@dataclass
class _Segment:
    start: int
    end: int
    startPadding: int
    endPadding: int

    @property
    def length(self) -> int:
        return self.end - self.start

    @property
    def padded_length(self) -> int:
        return self.end + self.endPadding - (self.start - self.startPadding)


def _exact_split_into_segments(length: int, exact: int, overlap: int) -> List[_Segment]:
    """
    Splits the given length into segments of `exact` (padded) length.
    Segments will overlap into each other with at least the given overlap.
    """
    if length == exact:
        # trivial
        return [_Segment(0, exact, 0, 0)]

    assert length > exact
    assert exact > overlap * 2

    result: List[_Segment] = []

    def add(s: _Segment):
        assert s.padded_length == exact
        result.append(s)

    # The current strategy is to go from left to right and to align segments
    # such that we use the least overlap possible. The last segment will then
    # be the smallest with potentially a lot of overlap.
    # While this is easy to implement, it's actually not ideal. Ideally, we
    # would want for the overlap to be distributed evenly between segments.
    # However, this is complex to implement and the current method also works.

    # we know that the first segment looks like this
    add(_Segment(0, exact - overlap, 0, overlap))

    while result[-1].end < length:
        startPadding = overlap
        start = result[-1].end
        end = start + exact - overlap * 2
        endPadding = overlap

        if end + endPadding >= length:
            # last segment
            endPadding = 0
            end = length
            startPadding = exact - (end - start)

        add(_Segment(start, end, startPadding, endPadding))

    return result


def _exact_split_into_regions(
    w: int,
    h: int,
    exact_w: int,
    exact_h: int,
    overlap: int,
) -> List[Tuple[Region, Padding]]:
    """
    Returns a list of disjoint regions along with padding.
    Each region plus its padding is guaranteed to have the given exact size.
    The padding (if not zero) is guaranteed to be at least the given overlap value.
    """

    # we can split x and y independently from each other and then combine the results
    x_segments = _exact_split_into_segments(w, exact_w, overlap)
    y_segments = _exact_split_into_segments(h, exact_h, overlap)

    logger.info(f"chaiNNer: image is split into {len(x_segments)}x{len(y_segments)} tiles each exactly {exact_w}x{exact_h}px")

    result: List[Tuple[Region, Padding]] = []
    for y in y_segments:
        for x in x_segments:
            result.append(
                (
                    Region(x.start, y.start, x.length, y.length),
                    Padding(y.startPadding, x.endPadding, y.endPadding, x.startPadding),
                )
            )
    return result


def _exact_split_without_padding(
    img: np.ndarray,
    exact_size: Size,
    upscale: Callable[[np.ndarray, Region], np.ndarray],
    overlap: int,
) -> np.ndarray:
    h, w, c = get_h_w_c(img)
    exact_w, exact_h = exact_size
    assert w >= exact_w and h >= exact_h

    if (w, h) == exact_size:
        return upscale(img, Region(0, 0, w, h))

    # To allocate the result image, we need to know the upscale factor first,
    # and we only get to know this factor after the first successful upscale.
    result: Optional[np.ndarray] = None
    scale: int = 0

    regions = _exact_split_into_regions(w, h, exact_w, exact_h, overlap)
    for tile, pad in regions:
        padded_tile = tile.add_padding(pad)

        upscale_result = upscale(padded_tile.read_from(img), padded_tile)

        # figure out by how much the image was upscaled by
        up_h, up_w, _ = get_h_w_c(upscale_result)
        current_scale = up_h // padded_tile.height
        assert current_scale > 0
        assert padded_tile.height * current_scale == up_h
        assert padded_tile.width * current_scale == up_w

        if result is None:
            # allocate the result image
            scale = current_scale
            result = np.zeros((h * scale, w * scale, c), dtype=np.float32)

        assert current_scale == scale

        # remove overlap padding
        upscale_result = pad.scale(scale).remove_from(upscale_result)

        # copy into result image
        tile.scale(scale).write_into(result, upscale_result)

    assert result is not None

    # remove initially added padding
    return result


def exact_split(
    img: np.ndarray,
    exact_size: Size,
    upscale: Callable[[np.ndarray, Region], np.ndarray],
    overlap: int = 16,
) -> np.ndarray:
    """
    Splits the image into tiles with exactly the given tile size.

    If the image is smaller than the given size, then it will be padded.
    """

    # ensure that the image is at least as large as the given size
    img, base_padding = _pad_image(img, exact_size)
    h, w, _ = get_h_w_c(img)

    result = _exact_split_without_padding(img, exact_size, upscale, overlap)
    scale = get_h_w_c(result)[0] // h

    if base_padding.empty:
        return result

    # remove initially added padding
    return (
        Region(0, 0, w, h).remove_padding(base_padding).scale(scale).read_from(result)
    )

nodes/impl/upscale/exact_split.py
vladmandic/sd-extension-chainner
from enum import Enum
from typing import List

import numpy as np

from ..color.convert import convert
from ..color.convert_data import LAB, RGB
from ..image_op import ImageOp


class SplitMode(Enum):
    RGB = 1
    LAB = 2

    def split(self, img: np.ndarray) -> List[np.ndarray]:
        if img.ndim == 2:
            return [img]

        assert img.ndim == 3
        c = img.shape[2]

        if c == 1:
            return [img[:, :, 0]]

        if self == SplitMode.RGB:
            return [img[:, :, channel] for channel in range(c)]
        elif self == SplitMode.LAB:
            if c < 3:
                return [img[:, :, channel] for channel in range(c)]
            lab = convert(img[:, :, 0:3], RGB, LAB)
            remaining_channels = [img[:, :, channel] for channel in range(3, c)]
            return [
                lab[:, :, 0],
                lab[:, :, 1],
                lab[:, :, 2],
                *remaining_channels,
            ]
        else:
            raise AssertionError

    def combine(self, channels: List[np.ndarray]) -> np.ndarray:
        l = len(channels)
        assert l > 0

        if l == 1:
            return channels[0]

        if self == SplitMode.RGB:
            return np.dstack(channels)
        elif self == SplitMode.LAB:
            if l < 3:
                return np.dstack(channels)
            rgb = convert(np.dstack(channels[0:3]), LAB, RGB)
            if l == 3:
                return rgb
            return np.dstack([rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2], *channels[3:]])
        else:
            raise AssertionError


def grayscale_split(
    img: np.ndarray, process: ImageOp, mode: SplitMode = SplitMode.RGB
) -> np.ndarray:
    """
    This function guarantees that the given image operation method will be called with 2D single-channel images.
    The images passed into the operation are guaranteed to have the same size as the given image.
    """

    input_channels = mode.split(img)
    output_channels: List[np.ndarray] = []
    for channel in input_channels:
        output_channels.append(process(channel))

    return mode.combine(output_channels)

nodes/impl/upscale/grayscale.py
vladmandic/sd-extension-chainner
import numpy as np

from ...utils.utils import get_h_w_c
from ..image_op import ImageOp


def passthrough_single_color(img: np.ndarray, scale: int, op: ImageOp) -> np.ndarray:
    """
    If the given image is a single-color image, it will be scaled and returned as is instead of being processed by the given operation.
    Obviously, this optimization is only correct if `op` doesn't change the color of single-color images.

    To make this a transparent optimization, it is important that `scale` is correct.
    `scale` must be the same factor by which `op` changes the dimension of the image.
    """

    h, w, c = get_h_w_c(img)

    if c == 1:
        unique_list = np.unique(img)
        if len(unique_list) == 1:
            return np.full((h * scale, w * scale), unique_list[0], np.float32)
    else:
        unique_values = []
        is_unique = True
        for channel in range(c):
            unique_list = np.unique(img[:, :, channel])
            if len(unique_list) == 1:
                unique_values.append(unique_list[0])
            else:
                is_unique = False
                break

        if is_unique:
            channels = [
                np.full((h * scale, w * scale), unique_values[channel], np.float32)
                for channel in range(c)
            ]
            return np.dstack(channels)

    return op(img)

nodes/impl/upscale/passthrough.py
vladmandic/sd-extension-chainner
from abc import ABC, abstractmethod

from ...utils.utils import Size


class Tiler(ABC):
    @abstractmethod
    def allow_smaller_tile_size(self) -> bool:
        """
        Whether the split implementation may use tile sizes smaller than the ones returned by this tiler.

        If False, then the split implementation guarantees that the all tiles are of exactly the size returned by this tiler.
        If the image is smaller than the returned tile size, then it has to be padded to reach the given size.
        """

    @abstractmethod
    def starting_tile_size(self, width: int, height: int, channels: int) -> Size:
        """
        The starting tile size is the first tile size that will be used.

        We generally prefer square tile sizes, but any tile size may be used.
        """

    def split(self, tile_size: Size) -> Size:
        w, h = tile_size
        assert w >= 16 and h >= 16
        return max(16, w // 2), max(16, h // 2)


class NoTiling(Tiler):
    def allow_smaller_tile_size(self) -> bool:
        return True

    def starting_tile_size(self, width: int, height: int, _channels: int) -> Size:
        size = max(width, height)
        # we prefer square tiles
        return size, size

    def split(self, _: Size) -> Size:
        raise ValueError("Image cannot be upscale with No Tiling mode.")


class MaxTileSize(Tiler):
    def __init__(self, tile_size: int = 2**31) -> None:
        self.tile_size: int = tile_size

    def allow_smaller_tile_size(self) -> bool:
        return True

    def starting_tile_size(self, width: int, height: int, _channels: int) -> Size:
        # Tile size a lot larger than the image don't make sense.
        # So we use the minimum of the image dimensions and the given tile size.
        max_tile_size = max(width + 10, height + 10)
        size = min(self.tile_size, max_tile_size)
        return size, size


class ExactTileSize(Tiler):
    def __init__(self, exact_size: Size) -> None:
        self.exact_size = exact_size

    def allow_smaller_tile_size(self) -> bool:
        return False

    def starting_tile_size(self, _width: int, _height: int, _channels: int) -> Size:
        return self.exact_size

    def split(self, _tile_size: Size) -> Size:
        raise ValueError(
            f"Splits are not supported for exact size ({self.exact_size[0]}x{self.exact_size[1]}px) splitting."
            f" This typically means that your machine does not have enough VRAM to run the current model."
        )

nodes/impl/upscale/tiler.py
vladmandic/sd-extension-chainner
import os
from typing import Tuple
import torch
from nodes.impl.pytorch.model_loading import load_state_dict
from nodes.impl.pytorch.types import PyTorchModel
from nodes.utils.unpickler import RestrictedUnpickle


def parse_ckpt_state_dict(checkpoint: dict):
    state_dict = {}
    for i, j in checkpoint.items():
        if "netG." in i:
            key = i.replace("netG.", "")
            state_dict[key] = j
        elif "module." in i:
            key = i.replace("module.", "")
            state_dict[key] = j
    return state_dict


def load_model(path: str, device, fp16: bool = False) -> Tuple[PyTorchModel, str, str]:
    """Read a pth file from the specified path and return it as a state dict
    and loaded model after finding arch config"""
    assert os.path.exists(path), f"Model file at location {path} does not exist"
    assert os.path.isfile(path), f"Path {path} is not a file"
    try:
        extension = os.path.splitext(path)[1].lower()

        if extension == ".pt":
            state_dict = torch.jit.load(  # type: ignore
                path, map_location=device
            ).state_dict()
        elif extension == ".pth":
            state_dict = torch.load(
                path,
                map_location=device,
                pickle_module=RestrictedUnpickle,  # type: ignore
            )
        elif extension == ".ckpt":
            checkpoint = torch.load(
                path,
                map_location=device,
                pickle_module=RestrictedUnpickle,  # type: ignore
            )
            if "state_dict" in checkpoint:
                checkpoint = checkpoint["state_dict"]
            state_dict = parse_ckpt_state_dict(checkpoint)
        else:
            raise ValueError(
                f"Unsupported model file extension {extension}. Please try a supported model type."
            )

        model = load_state_dict(state_dict)

        for _, v in model.named_parameters():
            v.requires_grad = False
        model.eval()
        model = model.to(device)
        should_use_fp16 = fp16
        if should_use_fp16:
            model = model.half()
        else:
            model = model.float()
    except Exception as e:
        raise ValueError(
            f"Model {os.path.basename(path)} is unsupported by chaiNNer. Please try"
            " another."
        ) from e

    return model

nodes/load_model.py
vladmandic/sd-extension-chainner
import logging
from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
from rich.theme import Theme
from rich.logging import RichHandler
from rich.console import Console
from rich.pretty import install

log = logging.getLogger("sd")
logger = log
console = Console(log_time=True, log_time_format='%H:%M:%S-%f', theme=Theme({
    "traceback.border": "black",
    "traceback.border.syntax_error": "black",
    "inspect.value.border": "black",
}))
install(console=console)

class PB(Progress):
    def __init__(self):
        super().__init__(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=console)


if not log.hasHandlers():
    log.setLevel(logging.DEBUG)
    rh = RichHandler(show_time=True, omit_repeated_times=False, show_level=True, show_path=False, markup=False, rich_tracebacks=True, log_time_format='%H:%M:%S-%f', level=logging.DEBUG, console=console)
    rh.setLevel(logging.DEBUG)
    logging.basicConfig(level=logging.ERROR, format='%(asctime)s %(levelname)s | %(module)s | %(message)s', handlers=[rh])
    log.debug('Initialized chaiNNer logging')

nodes/log.py
vladmandic/sd-extension-chainner
import functools
import hashlib
import os
import tempfile
import time
from enum import Enum
from typing import Dict, List, Optional, Tuple

import numpy as np
from nodes.log import logger

CACHE_MAX_BYTES = int(os.environ.get("CACHE_MAX_BYTES", 1024**3))  # default 1 GiB
CACHE_REGISTRY: List["NodeOutputCache"] = []


class CachedNumpyArray:
    def __init__(self, arr: np.ndarray):
        self.file = tempfile.TemporaryFile()
        self.file.write(arr.tobytes())

        self.shape = arr.shape
        self.dtype = arr.dtype

    def value(self) -> np.ndarray:
        self.file.seek(0)
        return np.frombuffer(self.file.read(), dtype=self.dtype).reshape(self.shape)


class NodeOutputCache:
    def __init__(self):
        self._data: Dict[Tuple, List] = {}
        self._bytes: Dict[Tuple, int] = {}
        self._access_time: Dict[Tuple, float] = {}

        CACHE_REGISTRY.append(self)

    @staticmethod
    def _args_to_key(args) -> Tuple:
        key = []
        for arg in args:
            if isinstance(arg, (int, float, bool, str, bytes)):
                key.append(arg)
            elif arg is None:
                key.append(None)
            elif isinstance(arg, Enum):
                key.append(arg.value)
            elif isinstance(arg, np.ndarray):
                key.append(tuple(arg.shape))
                key.append(arg.dtype.str)
                key.append(hashlib.sha256(arg.tobytes()).digest())
            elif hasattr(arg, "cache_key_func"):
                key.append(arg.__class__.__name__)
                key.append(arg.cache_key_func())
            else:
                raise RuntimeError(f"Unexpected argument type {arg.__class__.__name__}")
        return tuple(key)

    @staticmethod
    def _estimate_bytes(output) -> int:
        size = 0
        for out in output:
            if isinstance(out, np.ndarray):
                size += out.nbytes
            else:
                # any other type but numpy arrays is probably negligible, but here's an overestimate to handle
                # pathological cases where someone has a pipeline with a million math nodes
                size += 1024  # 1 KiB
        return size

    def empty(self):
        return len(self._data) == 0

    def oldest(self) -> Tuple[Tuple, float]:
        return min(self._access_time.items(), key=lambda x: x[1])

    def size(self):
        return sum(self._bytes.values())

    @staticmethod
    def _enforce_limits():
        while True:
            total_bytes = sum([cache.size() for cache in CACHE_REGISTRY])
            logger.debug(
                f"Cache size: {total_bytes} ({100*total_bytes/CACHE_MAX_BYTES:0.1f}% of limit)"
            )
            if total_bytes <= CACHE_MAX_BYTES:
                return
            logger.debug("Dropping oldest cache key")

            oldest_keys = [
                (cache, cache.oldest()) for cache in CACHE_REGISTRY if not cache.empty()
            ]

            cache, (key, _) = min(oldest_keys, key=lambda x: x[1][1])
            cache.drop(key)

    @staticmethod
    def _write_arrays_to_disk(output: List) -> List:
        return [
            CachedNumpyArray(item) if isinstance(item, np.ndarray) else item
            for item in output
        ]

    @staticmethod
    def _read_arrays_from_disk(output: List) -> List:
        return [
            item.value() if isinstance(item, CachedNumpyArray) else item
            for item in output
        ]

    @staticmethod
    def _output_to_list(output) -> List:
        if isinstance(output, list):
            return output
        elif isinstance(output, tuple):
            return list(output)
        else:
            return [output]

    @staticmethod
    def _list_to_output(output: List):
        if len(output) == 1:
            return output[0]
        return output

    def get(self, args) -> Optional[List]:
        key = self._args_to_key(args)
        if key in self._data:
            logger.debug("Cache hit")
            self._access_time[key] = time.time()
            return self._list_to_output(self._read_arrays_from_disk(self._data[key]))
        logger.debug("Cache miss")
        return None

    def put(self, args, output):
        key = self._args_to_key(args)
        self._data[key] = self._write_arrays_to_disk(self._output_to_list(output))
        self._bytes[key] = self._estimate_bytes(output)
        self._access_time[key] = time.time()
        self._enforce_limits()

    def drop(self, key):
        del self._data[key]
        del self._bytes[key]
        del self._access_time[key]


def cached(run):
    cache = NodeOutputCache()

    @functools.wraps(run)
    def _run(*args):
        out = cache.get(args)
        if out is not None:
            return out
        output = run(*args)
        cache.put(args, output)
        return output

    return _run

nodes/node_cache.py
vladmandic/sd-extension-chainner
from nodes.base_input import *

from .file_inputs import *
from .generic_inputs import *
from .image_dropdown_inputs import *
from .numeric_inputs import *
from .numpy_inputs import *

try:
    from .ncnn_inputs import *
except:
    pass
try:
    from .onnx_inputs import *
except:
    pass
try:
    from .pytorch_inputs import *
except:
    pass

nodes/properties/inputs/__init__.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import os
from typing import Literal, Union

from nodes.base_input import BaseInput

# pylint: disable=relative-beyond-top-level
from ...impl.image_formats import get_available_image_formats

FileInputKind = Union[
    Literal["bin"],
    Literal["image"],
    Literal["onnx"],
    Literal["param"],
    Literal["pt"],
    Literal["pth"],
    Literal["video"],
]


class FileInput(BaseInput):
    """Input for submitting a local file"""

    def __init__(
        self,
        input_type_name: str,
        label: str,
        file_kind: FileInputKind,
        filetypes: list[str],
        has_handle: bool = False,
        primary_input: bool = False,
    ):
        super().__init__(input_type_name, label, kind="file", has_handle=has_handle)
        self.filetypes = filetypes
        self.file_kind = file_kind
        self.primary_input = primary_input

        self.input_adapt = f"""
            match Input {{
                string as path => {input_type_name} {{ path: path }},
                _ => never
            }}
        """

        self.associated_type = str

    def toDict(self):
        return {
            **super().toDict(),
            "filetypes": self.filetypes,
            "fileKind": self.file_kind,
            "primaryInput": self.primary_input,
        }

    def enforce(self, value) -> str:
        assert isinstance(value, str)
        assert os.path.exists(value), f"File {value} does not exist"
        assert os.path.isfile(value), f"The path {value} is not a file"
        return value


def ImageFileInput(primary_input: bool = False) -> FileInput:
    """Input for submitting a local image file"""
    return FileInput(
        input_type_name="ImageFile",
        label="Image File",
        file_kind="image",
        filetypes=get_available_image_formats(),
        has_handle=False,
        primary_input=primary_input,
    )


def VideoFileInput(primary_input: bool = False) -> FileInput:
    """Input for submitting a local video file"""
    return FileInput(
        input_type_name="VideoFile",
        label="Video File",
        file_kind="video",
        filetypes=[
            ".mp4",
            ".h264",
            ".hevc",
            ".webm",
            ".avi",
            ".gif",
            ".mov",
            ".mkv",
            ".flv",
            ".m4v",
            ".avs",
        ],
        has_handle=False,
        primary_input=primary_input,
    )


def PthFileInput(primary_input: bool = False) -> FileInput:
    """Input for submitting a local .pth file"""
    return FileInput(
        input_type_name="PthFile",
        label="Model",
        file_kind="pth",
        filetypes=[".pt", ".pth", ".ckpt"],
        primary_input=primary_input,
    )


class DirectoryInput(BaseInput):
    """Input for submitting a local directory"""

    def __init__(
        self,
        label: str = "Base Directory",
        has_handle: bool = False,
        must_exist: bool = True,
        hide_label: bool = False,
    ):
        super().__init__("Directory", label, kind="directory", has_handle=has_handle)

        self.input_adapt = """
            match Input {
                string as path => Directory { path: path },
                _ => never
            }
        """

        self.must_exist: bool = must_exist
        self.hide_label: bool = hide_label

        self.associated_type = str

    def toDict(self):
        return {
            **super().toDict(),
            "hideLabel": self.hide_label,
        }

    def enforce(self, value):
        assert isinstance(value, str)
        if self.must_exist:
            assert os.path.exists(value), f"Directory {value} does not exist"
        return value


def BinFileInput(primary_input: bool = False) -> FileInput:
    """Input for submitting a local .bin file"""
    return FileInput(
        input_type_name="NcnnBinFile",
        label="NCNN Bin File",
        file_kind="bin",
        filetypes=[".bin"],
        primary_input=primary_input,
    )


def ParamFileInput(primary_input: bool = False) -> FileInput:
    """Input for submitting a local .param file"""
    return FileInput(
        input_type_name="NcnnParamFile",
        label="NCNN Param File",
        file_kind="param",
        filetypes=[".param"],
        primary_input=primary_input,
    )


def OnnxFileInput(primary_input: bool = False) -> FileInput:
    """Input for submitting a local .onnx file"""
    return FileInput(
        input_type_name="OnnxFile",
        label="ONNX Model File",
        file_kind="onnx",
        filetypes=[".onnx"],
        primary_input=primary_input,
    )

nodes/properties/inputs/file_inputs.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import json
import re
from enum import Enum
from typing import Dict, Generic, List, Literal, Tuple, Type, TypedDict, TypeVar, Union

import numpy as np
from nodes.log import logger

import navi
from nodes.base_input import BaseInput, InputConversion

from ...impl.blend import BlendMode
from ...impl.color.color import Color
from ...impl.dds.format import DDSFormat
from ...impl.image_utils import FillColor, normalize
from ...impl.upscale.auto_split_tiles import TileSize
from ...utils.format import format_color_with_channels
from ...utils.seed import Seed
from ...utils.utils import (
    join_pascal_case,
    join_space_case,
    split_pascal_case,
    split_snake_case,
)
from .numeric_inputs import NumberInput


class UntypedOption(TypedDict):
    option: str
    value: str | int


class TypedOption(TypedDict):
    option: str
    value: str | int
    type: navi.ExpressionJson


DropDownOption = Union[UntypedOption, TypedOption]

DropDownStyle = Literal["dropdown", "checkbox"]
"""
This specified the preferred style in which the frontend may display the dropdown.

- `dropdown`: This is the default style. The dropdown will simply be displayed as a dropdown.
- `checkbox`: If the dropdown has 2 options, then it will be displayed as a checkbox.
  The first option will be interpreted as the yes/true option while the second option will be interpreted as the no/false option.
"""


class DropDownInput(BaseInput):
    """Input for a dropdown"""

    def __init__(
        self,
        input_type: navi.ExpressionJson,
        label: str,
        options: List[DropDownOption],
        default_value: str | int | None = None,
        preferred_style: DropDownStyle = "dropdown",
        associated_type: Union[Type, None] = None,
    ):
        super().__init__(input_type, label, kind="dropdown", has_handle=False)
        self.options = options
        self.accepted_values = {o["value"] for o in self.options}
        self.default = (
            default_value if default_value is not None else options[0]["value"]
        )
        self.preferred_style: DropDownStyle = preferred_style

        if self.default not in self.accepted_values:
            logger.error(f"chaiNNer: invalid default value {self.default} in {label} dropdown. Using first value instead.")
            self.default = options[0]["value"]

        self.associated_type = (
            associated_type if associated_type is not None else type(self.default)
        )

    def toDict(self):
        return {
            **super().toDict(),
            "options": self.options,
            "def": self.default,
            "preferredStyle": self.preferred_style,
        }

    def make_optional(self):
        raise ValueError("DropDownInput cannot be made optional")

    def enforce(self, value):
        assert value in self.accepted_values, f"{value} is not a valid option"
        return value


class BoolInput(DropDownInput):
    def __init__(self, label: str, default: bool = True):
        super().__init__(
            input_type="bool",
            label=label,
            default_value=int(default),
            options=[
                {
                    "option": "Yes",
                    "value": int(True),  # 1
                    "type": "true",
                },
                {
                    "option": "No",
                    "value": int(False),  # 0
                    "type": "false",
                },
            ],
            preferred_style="checkbox",
        )
        self.associated_type = bool

    def enforce(self, value) -> bool:
        value = super().enforce(value)
        return bool(value)


T = TypeVar("T", bound=Enum)


class EnumInput(Generic[T], DropDownInput):
    """
    This adapts a python Enum into a chaiNNer dropdown input.

    ### Features

    All variants of the enum will be converted into typed dropdown options.
    The dropdown will be fully typed and bring its own type definitions.
    Option labels can be (partially) overridden using `option_labels`.

    By default, the input label, type names, and option labels will all be generated from the enum name and variant names.
    All of those defaults can be overridden.

    Options will be ordered by declaration order in the python enum definition.

    ### Requirements

    The value of each variant has to be either `str` or `int`.
    Other types are not permitted.
    """

    def __init__(
        self,
        enum: Type[T],
        label: str | None = None,
        default: T | None = None,
        type_name: str | None = None,
        option_labels: Dict[T, str] | None = None,
        extra_definitions: str | None = None,
    ):
        if type_name is None:
            type_name = enum.__name__
        if label is None:
            label = join_space_case(split_pascal_case(type_name))
        if option_labels is None:
            option_labels = {}

        options: List[DropDownOption] = []
        variant_types: List[str] = []
        for variant in enum:
            value = variant.value
            assert isinstance(value, (int, str))
            assert (
                re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", variant.name) is not None
            ), f"Expected the name of {enum.__name__}.{variant.name} to be snake case."

            name = split_snake_case(variant.name)
            variant_type = f"{type_name}::{join_pascal_case(name)}"
            option_label = option_labels.get(variant, join_space_case(name))

            variant_types.append(variant_type)
            options.append(
                {"option": option_label, "value": value, "type": variant_type}
            )

        super().__init__(
            input_type=type_name,
            label=label,
            options=options,
            default_value=default.value if default is not None else None,
        )

        self.type_definitions = (
            f"let {type_name} = {' | '.join(variant_types)};\n"
            + "\n".join([f"struct {t};" for t in variant_types])
            + (extra_definitions or "")
        )
        self.type_name: str = type_name
        self.enum = enum

        self.associated_type = enum

    def enforce(self, value) -> T:
        value = super().enforce(value)
        return self.enum(value)


class TextInput(BaseInput):
    """Input for arbitrary text"""

    def __init__(
        self,
        label: str,
        has_handle=True,
        min_length: int = 0,
        max_length: Union[int, None] = None,
        placeholder: Union[str, None] = None,
        multiline: bool = False,
        allow_numbers: bool = True,
        default: Union[str, None] = None,
        hide_label: bool = False,
        allow_empty_string: bool = False,
    ):
        super().__init__(
            input_type="string" if min_length == 0 else 'invStrSet("")',
            label=label,
            has_handle=has_handle,
            kind="text",
        )
        self.min_length = min_length
        self.max_length = max_length
        self.placeholder = placeholder
        self.default = default
        self.multiline = multiline
        self.hide_label = hide_label
        self.allow_empty_string = allow_empty_string

        if default is not None:
            assert default != ""
            assert min_length < len(default)
            assert max_length is None or len(default) < max_length

        self.associated_type = str

        if allow_numbers:
            self.input_conversions = [InputConversion("number", "toString(Input)")]

    def enforce(self, value) -> str:
        if isinstance(value, float) and int(value) == value:
            # stringify integers values
            value = str(int(value))
        else:
            value = str(value)

        # enforce length range
        if self.max_length is not None and len(value) > self.max_length:
            value = value[: self.max_length]
        if len(value) < self.min_length:
            raise ValueError(
                f"Text value of input '{self.label}' must be at least {self.min_length} characters long,"
                f" but found {len(value)} ('{value}')."
            )

        return value

    def toDict(self):
        return {
            **super().toDict(),
            "minLength": self.min_length,
            "maxLength": self.max_length,
            "placeholder": self.placeholder,
            "multiline": self.multiline,
            "def": self.default,
            "hideLabel": self.hide_label,
            "allowEmptyString": self.allow_empty_string,
        }


class ClipboardInput(BaseInput):
    """Input for pasting from clipboard"""

    def __init__(self, label: str = "Clipboard input"):
        super().__init__(["Image", "string", "number"], label, kind="text")
        self.input_conversions = [InputConversion("Image", '"<Image>"')]

    def enforce(self, value):
        if isinstance(value, np.ndarray):
            return normalize(value)

        if isinstance(value, float) and int(value) == value:
            # stringify integers values
            return str(int(value))

        return str(value)


class AnyInput(BaseInput):
    def __init__(self, label: str):
        super().__init__(input_type="any", label=label)
        self.associated_type = object

    def enforce_(self, value):
        # The behavior for optional inputs and None makes sense for all inputs except this one.
        return value


class SeedInput(NumberInput):
    def __init__(self, label: str = "Seed", has_handle: bool = True):
        super().__init__(
            label=label,
            minimum=None,
            maximum=None,
            precision=0,
            default=0,
        )
        self.has_handle = has_handle

        self.input_type = "Seed | int"
        self.input_conversions = [InputConversion("int", "Seed")]
        self.input_adapt = """
            match Input {
                int => Seed,
                _ => never
            }
        """

        self.associated_type = Seed

    def enforce(self, value) -> Seed:
        if isinstance(value, Seed):
            return value
        return Seed(int(value))

    def make_optional(self):
        raise ValueError("SeedInput cannot be made optional")


class ColorInput(BaseInput):
    def __init__(
        self,
        label: str = "Color",
        default: Color | None = None,
        channels: int | List[int] | None = None,
    ):
        super().__init__(
            input_type=navi.Color(channels=channels),
            label=label,
            has_handle=True,
            kind="color",
        )

        self.input_adapt = """
            match Input {
                string => parseColorJson(Input),
                _ => never
            }
        """

        self.channels: List[int] | None = (
            [channels] if isinstance(channels, int) else channels
        )

        if self.channels is None:
            if default is None:
                default = Color.bgr((0.5, 0.5, 0.5))
        else:
            assert len(self.channels) >= 0
            if default is None:
                if 3 in self.channels:
                    default = Color.bgr((0.5, 0.5, 0.5))
                elif 4 in self.channels:
                    default = Color.bgra((0.5, 0.5, 0.5, 1))
                elif 1 in self.channels:
                    default = Color.gray(0.5)
                else:
                    raise ValueError("Cannot find default color value")
            else:
                assert (
                    default.channels in self.channels
                ), "The default color is not accepted."

        self.default: Color = default

        self.associated_type = Color

    def enforce(self, value) -> Color:
        if isinstance(value, str):
            # decode color JSON strings from the frontend
            value = Color.from_json(json.loads(value))

        assert isinstance(value, Color)

        if self.channels is not None and value.channels not in self.channels:
            expected = format_color_with_channels(self.channels, plural=True)
            actual = format_color_with_channels([value.channels])
            raise ValueError(
                f"The input {self.label} only supports {expected} but was given {actual}."
            )

        return value

    def toDict(self):
        return {
            **super().toDict(),
            "def": json.dumps(self.default.to_json()),
            "channels": self.channels,
        }

    def make_optional(self):
        raise ValueError("ColorInput cannot be made optional")


def IteratorInput():
    """Input for showing that an iterator automatically handles the input"""
    return BaseInput("IteratorAuto", "Auto (Iterator)", has_handle=False)


class VideoContainer(Enum):
    MKV = "mkv"
    MP4 = "mp4"
    MOV = "mov"
    WEBM = "webm"
    AVI = "avi"
    GIF = "gif"
    NONE = "none"


VIDEO_CONTAINERS = {
    VideoContainer.MKV: "mkv",
    VideoContainer.MP4: "mp4",
    VideoContainer.MOV: "mov",
    VideoContainer.WEBM: "WebM",
    VideoContainer.AVI: "avi",
    VideoContainer.GIF: "GIF",
    VideoContainer.NONE: "None",
}


VIDEO_NONE_CONTAINERS: List[VideoContainer] = [VideoContainer.NONE, VideoContainer.GIF]


def VideoNoneContainerDropdown() -> DropDownInput:
    return DropDownInput(
        input_type="VideoContainer",
        label="Container",
        options=[
            {"option": VIDEO_CONTAINERS[vc], "value": vc.value}
            for vc in VIDEO_NONE_CONTAINERS
        ],
        associated_type=VideoContainer,
    )


VIDEO_FFV1_CONTAINERS: List[VideoContainer] = [VideoContainer.MKV]


def VideoFfv1ContainerDropdown() -> DropDownInput:
    return DropDownInput(
        input_type="VideoContainer",
        label="Container",
        options=[
            {"option": VIDEO_CONTAINERS[vc], "value": vc.value}
            for vc in VIDEO_FFV1_CONTAINERS
        ],
        associated_type=VideoContainer,
    )


VIDEO_VP9_CONTAINERS: List[VideoContainer] = [
    VideoContainer.WEBM,
    VideoContainer.MP4,
    VideoContainer.MKV,
]


def VideoVp9ContainerDropdown() -> DropDownInput:
    return DropDownInput(
        input_type="VideoContainer",
        label="Container",
        options=[
            {"option": VIDEO_CONTAINERS[vc], "value": vc.value}
            for vc in VIDEO_VP9_CONTAINERS
        ],
        associated_type=VideoContainer,
    )


VIDEO_H264_CONTAINERS: List[VideoContainer] = [
    VideoContainer.MKV,
    VideoContainer.MP4,
    VideoContainer.MOV,
    VideoContainer.AVI,
]


def VideoH264ContainerDropdown() -> DropDownInput:
    return DropDownInput(
        input_type="VideoContainer",
        label="Container",
        options=[
            {"option": VIDEO_CONTAINERS[vc], "value": vc.value}
            for vc in VIDEO_H264_CONTAINERS
        ],
        associated_type=VideoContainer,
    )


VIDEO_H265_CONTAINERS: List[VideoContainer] = [
    VideoContainer.MKV,
    VideoContainer.MP4,
    VideoContainer.MOV,
]


def VideoH265ContainerDropdown() -> DropDownInput:
    return DropDownInput(
        input_type="VideoContainer",
        label="Container",
        options=[
            {"option": VIDEO_CONTAINERS[vc], "value": vc.value}
            for vc in VIDEO_H265_CONTAINERS
        ],
        associated_type=VideoContainer,
    )


class VideoEncoder(Enum):
    H264 = "libx264"
    H265 = "libx265"
    VP9 = "libvpx-vp9"
    FFV1 = "ffv1"
    NONE = "none"


VIDEO_ENCODER_LABELS = {
    VideoEncoder.H264: "H.264 (AVC)",
    VideoEncoder.H265: "H.265 (HEVC)",
    VideoEncoder.VP9: "VP9",
    VideoEncoder.FFV1: "FFV1",
    VideoEncoder.NONE: "None",
}


def VideoEncoderDropdown() -> DropDownInput:
    return DropDownInput(
        input_type="VideoEncoder",
        label="Encoder",
        options=[
            {"option": label, "value": vc.value}
            for vc, label in VIDEO_ENCODER_LABELS.items()
        ],
        default_value=VideoEncoder.H264.value,
        associated_type=VideoEncoder,
    )


def VideoPresetDropdown() -> DropDownInput:
    """Video Type option dropdown"""
    return DropDownInput(
        input_type="VideoPreset",
        label="Preset",
        options=[
            {"option": "ultrafast", "value": "ultrafast"},
            {"option": "superfast", "value": "superfast"},
            {"option": "veryfast", "value": "veryfast"},
            {"option": "fast", "value": "fast"},
            {"option": "medium", "value": "medium"},
            {"option": "slow", "value": "slow"},
            {"option": "slower", "value": "slower"},
            {"option": "veryslow", "value": "veryslow"},
        ],
    )


def BlendModeDropdown() -> DropDownInput:
    """Blending Mode option dropdown"""
    return EnumInput(
        BlendMode,
        option_labels={
            BlendMode.ADD: "Linear Dodge (Add)",
        },
    )


def FillColorDropdown() -> DropDownInput:
    return EnumInput(
        FillColor,
        label="Negative Space Fill",
        default=FillColor.AUTO,
        extra_definitions="""
            def FillColor::getOutputChannels(fill: FillColor, channels: uint) {
                match fill {
                    FillColor::Transparent => 4,
                    _ => channels
                }
            }
        """,
    )


def TileSizeDropdown(
    label="Tile Size", estimate=True, default: TileSize | None = None
) -> DropDownInput:
    options = []
    if estimate:
        options.append({"option": "Auto (estimate)", "value": 0})

    options.append({"option": "Maximum", "value": -2})
    options.append({"option": "No Tiling", "value": -1})

    for size in [128, 192, 256, 384, 512, 768, 1024, 2048, 4096]:
        options.append({"option": str(size), "value": size})

    return DropDownInput(
        input_type="TileSize",
        label=label,
        options=options,
        associated_type=TileSize,
        default_value=default,
    )


SUPPORTED_DDS_FORMATS: List[Tuple[DDSFormat, str]] = [
    ("BC1_UNORM_SRGB", "BC1 (4bpp, sRGB, 1-bit Alpha)"),
    ("BC1_UNORM", "BC1 (4bpp, Linear, 1-bit Alpha)"),
    ("BC3_UNORM_SRGB", "BC3 (8bpp, sRGB, 8-bit Alpha)"),
    ("BC3_UNORM", "BC3 (8bpp, Linear, 8-bit Alpha)"),
    ("BC4_UNORM", "BC4 (4bpp, Grayscale)"),
    ("BC5_UNORM", "BC5 (8bpp, Unsigned, 2-channel normal)"),
    ("BC5_SNORM", "BC5 (8bpp, Signed, 2-channel normal)"),
    ("BC7_UNORM_SRGB", "BC7 (8bpp, sRGB, 8-bit Alpha)"),
    ("BC7_UNORM", "BC7 (8bpp, Linear, 8-bit Alpha)"),
    ("DXT1", "DXT1 (4bpp, Linear, 1-bit Alpha, Legacy)"),
    ("DXT3", "DXT3 (8bpp, Linear, 4-bit Alpha, Legacy)"),
    ("DXT5", "DXT5 (8bpp, Linear, 8-bit Alpha, Legacy)"),
    ("R8G8B8A8_UNORM_SRGB", "RGBA (32bpp, sRGB, 8-bit Alpha)"),
    ("R8G8B8A8_UNORM", "RGBA (32bpp, Linear, 8-bit Alpha)"),
    ("B8G8R8A8_UNORM_SRGB", "BGRA (32bpp, sRGB, 8-bit Alpha)"),
    ("B8G8R8A8_UNORM", "BGRA (32bpp, Linear, 8-bit Alpha)"),
    ("B5G5R5A1_UNORM", "BGRA (16bpp, Linear, 1-bit Alpha)"),
    ("B5G6R5_UNORM", "BGR (16bpp, Linear)"),
    ("B8G8R8X8_UNORM_SRGB", "BGRX (32bpp, sRGB)"),
    ("B8G8R8X8_UNORM", "BGRX (32bpp, Linear)"),
    ("R8G8_UNORM", "RG (16bpp, Linear)"),
    ("R8_UNORM", "R (8bpp, Linear)"),
]


def DdsFormatDropdown() -> DropDownInput:
    return DropDownInput(
        input_type="DdsFormat",
        label="DDS Format",
        options=[{"option": title, "value": f} for f, title in SUPPORTED_DDS_FORMATS],
        associated_type=DDSFormat,
    )


def DdsMipMapsDropdown() -> DropDownInput:
    return DropDownInput(
        input_type="DdsMipMaps",
        label="Generate Mip Maps",
        preferred_style="checkbox",
        options=[
            # these are not boolean values, see dds.py for more info
            {"option": "Yes", "value": 0},
            {"option": "No", "value": 1},
        ],
    )

nodes/properties/inputs/generic_inputs.py
vladmandic/sd-extension-chainner
import navi

from ...impl.color.convert_data import (
    color_spaces,
    color_spaces_or_detectors,
    get_alpha_partner,
    is_alpha_partner,
)

# pylint: disable=relative-beyond-top-level
from ...impl.image_utils import BorderType
from ...impl.pil_utils import InterpolationMethod, RotationInterpolationMethod
from .generic_inputs import DropDownInput, EnumInput


def ColorSpaceDetectorInput(label: str = "Color Space") -> DropDownInput:
    return DropDownInput(
        input_type="ColorSpace",
        label=label,
        options=[
            {
                "option": c.name,
                "value": c.id,
                "type": navi.named("ColorSpace", {"channels": c.channels}),
            }
            for c in color_spaces_or_detectors
        ],
    )


def ColorSpaceInput(label: str = "Color Space") -> DropDownInput:
    return DropDownInput(
        input_type="ColorSpace",
        label=label,
        options=[
            {
                "option": c.name,
                "value": c.id,
                "type": navi.named(
                    "ColorSpace",
                    {
                        "channels": c.channels,
                        "supportsAlpha": get_alpha_partner(c) is not None,
                    },
                ),
            }
            for c in color_spaces
            if not is_alpha_partner(c)
        ],
    )


def InterpolationInput() -> DropDownInput:
    """Resize interpolation dropdown"""
    return EnumInput(
        InterpolationMethod,
        option_labels={
            InterpolationMethod.NEAREST: "Nearest Neighbor",
            InterpolationMethod.BOX: "Area (Box)",
        },
    )


def RotateInterpolationInput() -> DropDownInput:
    return EnumInput(
        RotationInterpolationMethod,
        label="Interpolation Method",
        option_labels={
            RotationInterpolationMethod.NEAREST: "Nearest Neighbor",
        },
    )


def BorderInput() -> DropDownInput:
    return EnumInput(
        BorderType,
        default=BorderType.REFLECT_MIRROR,
        option_labels={
            BorderType.REFLECT_MIRROR: "Reflect (Mirror)",
            BorderType.WRAP: "Wrap (Tile)",
            BorderType.REPLICATE: "Replicate Edges",
        },
        extra_definitions="""
            def BorderType::getOutputChannels(type: BorderType, channels: uint, color: Color | null): uint {
                match type {
                    BorderType::Transparent => 4,
                    BorderType::CustomColor => match color {
                        Color => max(color.channels, channels),
                        null => never,
                    },
                    _ => channels
                }
            }
        """,
    )


def NormalChannelInvertInput() -> DropDownInput:
    return DropDownInput(
        input_type="NormalChannelInvert",
        label="Invert",
        options=[
            {
                "option": "None",
                "value": 0,
            },
            {
                "option": "Invert R",
                "value": 1,
            },
            {
                "option": "Invert G",
                "value": 2,
            },
            {
                "option": "Invert R and G",
                "value": 3,
            },
        ],
    )

nodes/properties/inputs/image_dropdown_inputs.py
vladmandic/sd-extension-chainner
from nodes.base_input import BaseInput

from ...impl.ncnn.model import NcnnModelWrapper


class NcnnModelInput(BaseInput):
    """Input for NcnnModel"""

    def __init__(self, label: str = "Model"):
        super().__init__("NcnnNetwork", label)
        self.associated_type = NcnnModelWrapper

nodes/properties/inputs/ncnn_inputs.py
vladmandic/sd-extension-chainner
import math
from typing import List, Literal, Tuple, Union

import navi
from nodes.base_input import BaseInput, InputConversion, InputKind

from ...utils.utils import round_half_up


def clampNumber(
    value: Union[float, int],
    precision: int,
    min_value: Union[float, int, None],
    max_value: Union[float, int, None],
) -> Union[float, int]:
    # Convert proper number type
    value = round_half_up(value) if precision == 0 else round(value, precision)

    # Clamp to max and min, correcting for max/min not aligning with offset + n * step
    if max_value is not None:
        value = min(value, max_value)
    if min_value is not None:
        value = max(value, min_value)

    # guarantee integers
    if precision <= 0:
        return int(value)
    else:
        return float(value)


def get_number_type(
    min_value: Union[float, int, None],
    max_value: Union[float, int, None],
    precision: int,
) -> navi.ExpressionJson:
    if precision > 0:
        # step is not an integer
        return navi.interval(min_value, max_value)
    return navi.int_interval(min_value, max_value)


class NumberInput(BaseInput):
    """Input a number"""

    def __init__(
        self,
        label: str,
        precision: int = 0,
        controls_step: Union[float, int, None] = None,
        default: Union[float, int] = 0,
        minimum: Union[float, int, None] = 0,
        maximum: Union[float, int, None] = None,
        unit: Union[str, None] = None,
        note_expression: Union[str, None] = None,
        kind: InputKind = "number",
        hide_trailing_zeros: bool = True,
        hide_label: bool = False,
    ):
        super().__init__("number", label, kind=kind, has_handle=True)
        self.precision = precision
        # controls_step is for increment/decrement arrows.
        self.controls_step: Union[float, int] = (
            controls_step if controls_step is not None else 10**-precision
        )
        self.default = default
        self.minimum = minimum
        self.maximum = maximum
        self.unit = unit
        self.note_expression = note_expression
        self.hide_trailing_zeros = hide_trailing_zeros
        self.hide_label = hide_label

        self.associated_type = float if precision > 0 else int

        self.input_type = get_number_type(
            self.minimum,
            self.maximum,
            self.precision,
        )
        if self.precision == 0:
            self.input_conversions = [InputConversion("number", "round(Input)")]

    def toDict(self):
        return {
            **super().toDict(),
            "min": self.minimum,
            "max": self.maximum,
            "noteExpression": self.note_expression,
            "def": self.default,
            "precision": self.precision,
            "controlsStep": self.controls_step,
            "unit": self.unit,
            "hideTrailingZeros": self.hide_trailing_zeros,
            "hideLabel": self.hide_label,
        }

    def make_optional(self):
        raise ValueError("NumberInput and SliderInput cannot be made optional")

    def enforce(self, value):
        assert isinstance(value, (int, float))

        if math.isnan(value):
            raise ValueError("NaN is not a valid number")

        return clampNumber(value, self.precision, self.minimum, self.maximum)


class SliderInput(NumberInput):
    """Input for integer number via slider"""

    def __init__(
        self,
        label: str,
        precision: int = 0,
        controls_step: Union[float, int, None] = None,
        slider_step: Union[float, int, None] = None,
        minimum: Union[float, int] = 0,
        maximum: Union[float, int] = 100,
        default: Union[float, int] = 50,
        unit: Union[str, None] = None,
        note_expression: Union[str, None] = None,
        ends: Tuple[Union[str, None], Union[str, None]] = (None, None),
        hide_trailing_zeros: bool = False,
        gradient: Union[List[str], None] = None,
        scale: Literal["linear", "log", "log-offset", "sqrt"] = "linear",
    ):
        super().__init__(
            label,
            precision=precision,
            controls_step=controls_step,
            default=default,
            minimum=minimum,
            maximum=maximum,
            unit=unit,
            note_expression=note_expression,
            kind="slider",
            hide_trailing_zeros=hide_trailing_zeros,
        )
        self.ends = ends
        self.slider_step = (
            slider_step
            if slider_step is not None
            else (controls_step if controls_step is not None else 10**-precision)
        )
        self.gradient = gradient
        self.scale = scale

    def toDict(self):
        return {
            **super().toDict(),
            "ends": self.ends,
            "sliderStep": self.slider_step,
            "gradient": self.gradient,
            "scale": self.scale,
        }

nodes/properties/inputs/numeric_inputs.py
vladmandic/sd-extension-chainner
# pylint: disable=relative-beyond-top-level
from __future__ import annotations

from typing import List, Optional, Union

import numpy as np

import navi
from nodes.base_input import BaseInput, ErrorValue

from ...impl.color.color import Color
from ...utils.format import format_color_with_channels, format_image_with_channels
from ...utils.utils import get_h_w_c


class AudioInput(BaseInput):
    """Input a 1D Audio NumPy array"""

    def __init__(self, label: str = "Audio"):
        super().__init__("Audio", label)


class ImageInput(BaseInput):
    """Input a 2D Image NumPy array"""

    def __init__(
        self,
        label: str = "Image",
        image_type: navi.ExpressionJson = "Image | Color",
        channels: Union[int, List[int], None] = None,
        allow_colors: bool = False,
    ):
        base_type = [navi.Image(channels=channels)]
        if allow_colors:
            base_type.append(navi.Color(channels=channels))
        image_type = navi.intersect(image_type, base_type)
        super().__init__(image_type, label)

        self.channels: Optional[List[int]] = (
            [channels] if isinstance(channels, int) else channels
        )
        self.allow_colors: bool = allow_colors

        self.associated_type = np.ndarray

        if self.allow_colors:
            self.associated_type = Union[np.ndarray, Color]

    def enforce(self, value):
        if isinstance(value, Color):
            if not self.allow_colors:
                raise ValueError(
                    f"The input {self.label} does not accept colors, but was connected with one."
                )

            if self.channels is not None and value.channels not in self.channels:
                expected = format_color_with_channels(self.channels, plural=True)
                actual = format_color_with_channels([value.channels])
                raise ValueError(
                    f"The input {self.label} only supports {expected} but was given {actual}."
                )

            return value

        assert isinstance(value, np.ndarray)
        _, _, c = get_h_w_c(value)

        if self.channels is not None and c not in self.channels:
            expected = format_image_with_channels(self.channels, plural=True)
            actual = format_image_with_channels([c])
            raise ValueError(
                f"The input {self.label} only supports {expected} but was given {actual}."
            )

        assert value.dtype == np.float32, "Expected the input image to be normalized."

        if c == 1 and value.ndim == 3:
            value = value[:, :, 0]

        return value

    def get_error_value(self, value) -> ErrorValue:
        def get_channels(channel: int) -> str:
            if channel == 1:
                return "Grayscale"
            if channel == 3:
                return "RGB"
            if channel == 4:
                return "RGBA"
            return f"{channel}-channel"

        if isinstance(value, Color):
            return {
                "type": "formatted",
                "formatString": f"{get_channels(value.channels)} Color",
            }
        elif isinstance(value, np.ndarray):
            h, w, c = get_h_w_c(value)
            return {
                "type": "formatted",
                "formatString": f"{get_channels(c)} Image {w}x{h}",
            }
        else:
            return super().get_error_value(value)


class VideoInput(BaseInput):
    """Input a 3D Video NumPy array"""

    def __init__(self, label: str = "Video"):
        super().__init__("Video", label)

nodes/properties/inputs/numpy_inputs.py
vladmandic/sd-extension-chainner
import navi
from nodes.base_input import BaseInput

from ...impl.onnx.model import OnnxModel, OnnxModels, OnnxRemBg, is_rembg_model
from .generic_inputs import DropDownInput


class OnnxModelInput(BaseInput):
    """Input for onnx model"""

    def __init__(
        self, label: str = "Model", input_type: navi.ExpressionJson = "OnnxModel"
    ):
        super().__init__(input_type, label)
        self.associated_type = OnnxModel


class OnnxGenericModelInput(OnnxModelInput):
    """ONNX model input for things that aren't background removal"""

    def __init__(
        self, label: str = "Model", input_type: navi.ExpressionJson = "OnnxModel"
    ):
        super().__init__(label, navi.intersect(input_type, "OnnxGenericModel"))

    def enforce(self, value):
        assert isinstance(value, OnnxModels)
        assert not is_rembg_model(value.bytes), "Expected a non-rembg model"
        return value


class OnnxRemBgModelInput(OnnxModelInput):
    """ONNX model input for background removal"""

    def __init__(
        self, label: str = "Model", input_type: navi.ExpressionJson = "OnnxModel"
    ):
        super().__init__(label, navi.intersect(input_type, "OnnxRemBgModel"))
        self.associated_type = OnnxRemBg

    def enforce(self, value):
        assert isinstance(value, OnnxModels)
        assert is_rembg_model(value.bytes), "Expected a rembg model"
        return value


def OnnxFpDropdown() -> DropDownInput:
    return DropDownInput(
        input_type="FpMode",
        label="Data Type",
        options=[
            {
                "option": "fp32",
                "value": 0,
                "type": "FpMode::fp32",
            },
            {
                "option": "fp16",
                "value": 1,
                "type": "FpMode::fp16",
            },
        ],
    )

nodes/properties/inputs/onnx_inputs.py
vladmandic/sd-extension-chainner
try:
    import torch

    from ...impl.pytorch.types import (
        PyTorchFaceModel,
        PyTorchInpaintModel,
        PyTorchModel,
        PyTorchSRModel,
        is_pytorch_face_model,
        is_pytorch_inpaint_model,
        is_pytorch_model,
        is_pytorch_sr_model,
    )
except:
    torch = None

import navi
from nodes.base_input import BaseInput


class ModelInput(BaseInput):
    """Input a loaded model"""

    def __init__(
        self,
        label: str = "Model",
        input_type: navi.ExpressionJson = "PyTorchModel",
    ):
        super().__init__(input_type, label)
        if torch is not None:
            self.associated_type = PyTorchModel

    def enforce(self, value):
        if torch is not None:
            assert isinstance(value, torch.nn.Module), "Expected a PyTorch model."
            assert is_pytorch_model(value), "Expected a supported PyTorch model."
        return value


class SrModelInput(ModelInput):
    def __init__(
        self,
        label: str = "Model",
        input_type: navi.ExpressionJson = "PyTorchModel",
    ):
        super().__init__(
            label,
            navi.intersect(input_type, "PyTorchSRModel"),
        )
        if torch is not None:
            self.associated_type = PyTorchSRModel

    def enforce(self, value):
        if torch is not None:
            assert isinstance(value, torch.nn.Module), "Expected a PyTorch model."
            assert is_pytorch_sr_model(
                value
            ), "Expected a regular Super-Resolution model."
        return value


class FaceModelInput(ModelInput):
    def __init__(
        self, label: str = "Model", input_type: navi.ExpressionJson = "PyTorchModel"
    ):
        super().__init__(
            label,
            navi.intersect(input_type, "PyTorchFaceModel"),
        )
        if torch is not None:
            self.associated_type = PyTorchFaceModel

    def enforce(self, value):
        if torch is not None:
            assert isinstance(value, torch.nn.Module), "Expected a PyTorch model."
            assert is_pytorch_face_model(
                value
            ), "Expected a Face-specific Super-Resolution model."
        return value


class InpaintModelInput(ModelInput):
    def __init__(
        self, label: str = "Model", input_type: navi.ExpressionJson = "PyTorchModel"
    ):
        super().__init__(
            label,
            navi.intersect(input_type, "PyTorchInpaintModel"),
        )
        if torch is not None:
            self.associated_type = PyTorchInpaintModel

    def enforce(self, value):
        if torch is not None:
            assert isinstance(value, torch.nn.Module), "Expected a PyTorch model."
            assert is_pytorch_inpaint_model(
                value
            ), "Expected an inpainting-specific model."
        return value


class TorchScriptInput(BaseInput):
    """Input a JIT traced model"""

    def __init__(self, label: str = "Traced Model"):
        super().__init__("PyTorchScript", label)

nodes/properties/inputs/pytorch_inputs.py
vladmandic/sd-extension-chainner
from nodes.base_output import *

from .file_outputs import *
from .generic_outputs import *
from .numpy_outputs import *

try:
    from .ncnn_outputs import *
except:
    pass
try:
    from .onnx_outputs import *
except:
    pass
try:
    from .pytorch_outputs import *
except:
    pass

nodes/properties/outputs/__init__.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import navi
from nodes.base_output import BaseOutput


class DirectoryOutput(BaseOutput):
    """Output for saving to a directory"""

    def __init__(
        self,
        label: str = "Directory",
        of_input: int | None = None,
        output_type: str = "Directory",
    ):
        directory_type = (
            "Directory"
            if of_input is None
            else f"splitFilePath(Input{of_input}.path).dir"
        )
        directory_type = navi.intersect(directory_type, output_type)
        super().__init__(directory_type, label, associated_type=str)

    def get_broadcast_type(self, value: str):
        return navi.named("Directory", {"path": navi.literal(value)})

    def enforce(self, value) -> str:
        assert isinstance(value, str)
        return value

nodes/properties/outputs/file_outputs.py
vladmandic/sd-extension-chainner
from __future__ import annotations

from typing import Union

import navi
from nodes.base_output import BaseOutput

from ...impl.color.color import Color
from ...utils.format import format_color_with_channels
from ...utils.seed import Seed


class NumberOutput(BaseOutput):
    def __init__(
        self,
        label: str,
        output_type: navi.ExpressionJson = "number",
    ):
        super().__init__(
            navi.intersect("number", output_type),
            label,
            associated_type=Union[int, float],
        )

    def get_broadcast_type(self, value: int | float):
        return navi.literal(value)

    def enforce(self, value) -> int | float:
        assert isinstance(value, (int, float))
        return value


class TextOutput(BaseOutput):
    def __init__(
        self,
        label: str,
        output_type: navi.ExpressionJson = "string",
    ):
        super().__init__(navi.intersect("string", output_type), label)

    def get_broadcast_type(self, value: str):
        return navi.literal(value)

    def enforce(self, value) -> str:
        assert isinstance(value, str)
        return value


def FileNameOutput(label: str = "Name", of_input: int | None = None):
    output_type = (
        "string"
        if of_input is None
        else f"splitFilePath(Input{of_input}.path).basename"
    )

    return TextOutput(label=label, output_type=output_type)


class SeedOutput(BaseOutput):
    def __init__(self, label: str = "Seed"):
        super().__init__(output_type="Seed", label=label, kind="generic")

    def enforce(self, value) -> Seed:
        assert isinstance(value, Seed)
        return value


class ColorOutput(BaseOutput):
    def __init__(
        self,
        label: str = "Color",
        color_type: navi.ExpressionJson = "Color",
        channels: int | None = None,
    ):
        super().__init__(
            output_type=navi.intersect(color_type, navi.Color(channels=channels)),
            label=label,
            kind="generic",
        )

        self.channels = channels

    def enforce(self, value) -> Color:
        assert isinstance(value, Color)

        if self.channels is not None and value.channels != self.channels:
            expected = format_color_with_channels([self.channels])
            actual = format_color_with_channels([value.channels])
            raise ValueError(
                f"The output {self.label} was supposed to return {expected} but actually returned {actual}."
                f" This is a bug in the implementation of the node."
                f" Please report this bug."
            )

        return value

nodes/properties/outputs/generic_outputs.py
vladmandic/sd-extension-chainner
import navi
from nodes.base_output import BaseOutput, OutputKind

from ...impl.ncnn.model import NcnnModelWrapper
from ...utils.format import format_channel_numbers


class NcnnModelOutput(BaseOutput):
    def __init__(
        self,
        model_type: navi.ExpressionJson = "NcnnNetwork",
        label: str = "Model",
        kind: OutputKind = "generic",
    ):
        super().__init__(model_type, label, kind=kind, associated_type=NcnnModelWrapper)

    def get_broadcast_data(self, value: NcnnModelWrapper):
        return {
            "tags": [
                format_channel_numbers(value.in_nc, value.out_nc),
                f"{value.nf}nf",
                value.fp,
            ]
        }

    def get_broadcast_type(self, value: NcnnModelWrapper):
        return navi.named(
            "NcnnNetwork",
            {
                "scale": value.scale,
                "inputChannels": value.in_nc,
                "outputChannels": value.out_nc,
                "nf": value.nf,
                "fp": navi.literal(value.fp),
            },
        )

nodes/properties/outputs/ncnn_outputs.py
vladmandic/sd-extension-chainner
import base64
from typing import Optional, Tuple

import cv2
import numpy as np

import navi
from nodes.base_output import BaseOutput, OutputKind

from ...impl.image_utils import normalize, to_uint8
from ...impl.pil_utils import InterpolationMethod, resize
from ...utils.format import format_image_with_channels
from ...utils.utils import get_h_w_c, round_half_up


class NumPyOutput(BaseOutput):
    """Output a NumPy array"""

    def __init__(
        self,
        output_type: navi.ExpressionJson,
        label: str,
        kind: OutputKind = "generic",
        has_handle: bool = True,
    ):
        super().__init__(
            output_type,
            label,
            kind=kind,
            has_handle=has_handle,
            associated_type=np.ndarray,
        )

    def enforce(self, value) -> np.ndarray:
        assert isinstance(value, np.ndarray)
        return value


def AudioOutput():
    """Output a 1D Audio NumPy array"""
    return NumPyOutput("Audio", "Audio")


class ImageOutput(NumPyOutput):
    def __init__(
        self,
        label: str = "Image",
        image_type: navi.ExpressionJson = "Image",
        kind: OutputKind = "image",
        has_handle: bool = True,
        channels: Optional[int] = None,
        assume_normalized: bool = False,
    ):
        super().__init__(
            navi.intersect(image_type, navi.Image(channels=channels)),
            label,
            kind=kind,
            has_handle=has_handle,
        )

        self.channels: Optional[int] = channels
        self.assume_normalized: bool = assume_normalized

    def get_broadcast_data(self, value: np.ndarray):
        h, w, c = get_h_w_c(value)
        return {
            "height": h,
            "width": w,
            "channels": c,
        }

    def get_broadcast_type(self, value: np.ndarray):
        h, w, c = get_h_w_c(value)
        return navi.Image(width=w, height=h, channels=c)

    def enforce(self, value) -> np.ndarray:
        assert isinstance(value, np.ndarray)

        h, w, c = get_h_w_c(value)

        if h == 0 or w == 0:
            raise ValueError(
                f"The output {self.label} returned an empty image (w={w} h={h})."
                f" This is a bug in the implementation of the node."
                f" Please report this bug."
            )

        if self.channels is not None and c != self.channels:
            expected = format_image_with_channels([self.channels])
            actual = format_image_with_channels([c])
            raise ValueError(
                f"The output {self.label} was supposed to return {expected} but actually returned {actual}."
                f" This is a bug in the implementation of the node."
                f" Please report this bug."
            )

        # flatting 3D single-channel images to 2D
        if c == 1 and value.ndim == 3:
            value = value[:, :, 0]

        if not self.assume_normalized:
            value = normalize(value)

        assert value.dtype == np.float32, (
            f"The output {self.label} did not return a normalized image."
            f" This is a bug in the implementation of the node."
            f" Please report this bug."
            f"\n\nTo the author of this node: Either use `normalize` or remove `assume_normalized=True` from this output."
        )

        # make image readonly
        value.setflags(write=False)

        return value


def preview_encode(
    img: np.ndarray,
    target_size: int = 512,
    grace: float = 1.2,
    lossless: bool = False,
) -> Tuple[str, np.ndarray]:
    """
    resize the image, so the preview loads faster and doesn't lag the UI
    512 was chosen as the default target because a 512x512 RGBA 8bit PNG is at most 1MB in size
    """
    h, w, c = get_h_w_c(img)

    max_size = target_size * grace
    if w > max_size or h > max_size:
        f = max(w / target_size, h / target_size)
        t = (max(1, round_half_up(w / f)), max(1, round_half_up(h / f)))
        if c == 4:
            # https://github.com/chaiNNer-org/chaiNNer/issues/1321
            img = resize(img, t, InterpolationMethod.BOX)
        else:
            img = cv2.resize(img, t, interpolation=cv2.INTER_AREA)

    image_format = "png" if c > 3 or lossless else "jpg"

    _, encoded_img = cv2.imencode(f".{image_format}", to_uint8(img, normalized=True))  # type: ignore
    base64_img = base64.b64encode(encoded_img).decode("utf8")  # type: ignore

    return f"data:image/{image_format};base64,{base64_img}", img


class LargeImageOutput(ImageOutput):
    def __init__(
        self,
        label: str = "Image",
        image_type: navi.ExpressionJson = "Image",
        kind: OutputKind = "large-image",
        has_handle: bool = True,
        assume_normalized: bool = False,
    ):
        super().__init__(
            label,
            image_type,
            kind=kind,
            has_handle=has_handle,
            assume_normalized=assume_normalized,
        )

    def get_broadcast_data(self, value: np.ndarray):
        img = value
        h, w, c = get_h_w_c(img)
        image_size = max(h, w)

        preview_sizes = [2048, 1024, 512, 256]
        preview_size_grace = 1.2

        start_index = len(preview_sizes) - 1
        for i, size in enumerate(preview_sizes):
            if size <= image_size and image_size <= size * preview_size_grace:
                # this preview size will perfectly fit the image
                start_index = i
                break
            if image_size > size:
                # the image size is larger than the preview size, so try to pick the previous size
                start_index = max(0, i - 1)
                break

        previews = []

        # Encode for multiple scales. Use the preceding scale to save time encoding the smaller sizes.
        last_encoded = img
        for size in preview_sizes[start_index:]:
            largest_preview = size == preview_sizes[start_index]
            url, last_encoded = preview_encode(
                last_encoded,
                target_size=size,
                grace=preview_size_grace,
                lossless=largest_preview,
            )
            le_h, le_w, _ = get_h_w_c(last_encoded)
            previews.append({"width": le_w, "height": le_h, "url": url})

        return {
            "previews": previews,
            "height": h,
            "width": w,
            "channels": c,
        }


def VideoOutput():
    """Output a 3D Video NumPy array"""
    return NumPyOutput("Video", "Video")

nodes/properties/outputs/numpy_outputs.py
vladmandic/sd-extension-chainner
import navi
from nodes.base_output import BaseOutput

from ...impl.onnx.model import OnnxModel


class OnnxModelOutput(BaseOutput):
    """Output for onnx model"""

    def __init__(
        self,
        model_type: navi.ExpressionJson = "OnnxModel",
        label: str = "Model",
    ):
        super().__init__(model_type, label, associated_type=OnnxModel)

    def get_broadcast_type(self, value: OnnxModel):
        fields = {
            "subType": navi.literal(value.sub_type),
        }

        if value.scale_width:
            fields["scaleWidth"] = value.scale_width
        if value.scale_height:
            fields["scaleHeight"] = value.scale_height

        return navi.named("OnnxModel", fields)

nodes/properties/outputs/onnx_outputs.py
vladmandic/sd-extension-chainner
from typing import List

import navi
from nodes.base_output import BaseOutput, OutputKind

from ...impl.pytorch.types import PyTorchModel
from ...utils.format import format_channel_numbers


def _get_sizes(value: PyTorchModel) -> List[str]:
    if "SRVGG" in value.model_arch:
        return [f"{value.num_feat}nf", f"{value.num_conv}nc"]
    elif (
        "SwinIR" in value.model_arch
        or "Swin2SR" in value.model_arch
        or "HAT" in value.model_arch
        or "SRFormer" in value.model_arch
    ):
        head_length = len(value.depths)  # type: ignore
        if head_length <= 4:
            size_tag = "small"
        elif head_length < 9:
            size_tag = "medium"
        else:
            size_tag = "large"
        return [
            size_tag,
            f"s{value.img_size}w{value.window_size}",
            f"{value.num_feat}nf",
            f"{value.embed_dim}dim",
            f"{value.resi_connection}",
        ]
    elif "DAT" in value.model_arch:
        head_length = len(value.depth)  # type: ignore
        if head_length <= 4:
            size_tag = "small"
        elif head_length < 9:
            size_tag = "medium"
        else:
            size_tag = "large"
        return [
            size_tag,
            f"s{value.img_size}|{value.split_size[0]}x{value.split_size[1]}",  # type: ignore
            f"{value.num_feat}nf",
            f"{value.embed_dim}dim",
            f"{value.resi_connection}",
        ]
    elif "OmniSR" in value.model_arch:
        return [
            f"{value.num_feat}nf",
            f"w{value.window_size}",
            f"{value.res_num}nr",
        ]
    elif value.model_arch in [
        "GFPGAN",
        "RestoreFormer",
        "CodeFormer",
        "LaMa",
        "MAT",
        "SCUNet",
    ]:
        return []
    else:
        return [
            f"{value.num_filters}nf",
            f"{value.num_blocks}nb",
        ]


class ModelOutput(BaseOutput):
    def __init__(
        self,
        model_type: navi.ExpressionJson = "PyTorchModel",
        label: str = "Model",
        kind: OutputKind = "generic",
    ):
        super().__init__(model_type, label, kind=kind, associated_type=PyTorchModel)

    def get_broadcast_data(self, value: PyTorchModel):
        return {
            "tags": [
                value.model_arch,
                format_channel_numbers(value.in_nc, value.out_nc),
                *_get_sizes(value),
            ]
        }

    def get_broadcast_type(self, value: PyTorchModel):
        return navi.named(
            "PyTorchModel",
            {
                "scale": value.scale,
                "inputChannels": value.in_nc,
                "outputChannels": value.out_nc,
                "arch": navi.literal(value.model_arch),
                "subType": navi.literal(value.sub_type),
                "size": navi.literal("x".join(_get_sizes(value))),
            },
        )


def TorchScriptOutput():
    """Output a JIT traced model"""
    return BaseOutput("PyTorchScript", "Traced Model")

nodes/properties/outputs/pytorch_outputs.py
vladmandic/sd-extension-chainner
from typing import Type, TypeVar

T = TypeVar("T")


def checked_cast(t: Type[T], value) -> T:
    assert isinstance(value, t), f"Value is {type(value)}, must be type {t}"
    return value

nodes/utils/checked_cast.py
vladmandic/sd-extension-chainner
from __future__ import annotations

from typing import Dict

from nodes.log import logger

SettingsJson = Dict[str, object]
JsonExecutionOptions = Dict[str, SettingsJson]


class ExecutionOptions:
    def __init__(
        self,
        backend_settings: JsonExecutionOptions,
    ) -> None:
        self.__settings = backend_settings

        logger.info(f"chaiNNer: execution options: {self.__settings}")

    @staticmethod
    def parse(json: JsonExecutionOptions) -> ExecutionOptions:
        return ExecutionOptions(backend_settings=json)

    def get_package_settings(self, package_id: str) -> SettingsJson:
        return self.__settings.get(package_id, {})


__global_exec_options = ExecutionOptions({})


def get_execution_options() -> ExecutionOptions:
    return __global_exec_options


def set_execution_options(value: ExecutionOptions):
    # TODO: Make the mutable global state unnecessary
    # pylint: disable=global-statement
    global __global_exec_options
    __global_exec_options = value

nodes/utils/exec_options.py
vladmandic/sd-extension-chainner
from typing import Callable, Iterable, List, Literal, TypeVar

T = TypeVar("T")
Conj = Literal["and", "or"]


def join_english(
    items: Iterable[T],
    to_str: Callable[[T], str] = str,
    conj: Conj = "and",
) -> str:
    s = list(map(to_str, items))

    l = len(s)
    assert l > 0

    if l == 1:
        return s[0]
    if l == 2:
        return f"{s[0]} {conj} {s[1]}"
    return ", ".join(s[:-1]) + f", {conj} " + s[-1]


def format_image_with_channels(
    channels: List[int],
    conj: Conj = "and",
    plural: bool = False,
) -> str:
    assert len(channels) > 0

    named = {1: "grayscale", 3: "RGB", 4: "RGBA"}
    if all(x in named for x in channels):
        if plural:
            return join_english(channels, lambda c: named[c], conj=conj) + " images"
        else:
            return (
                "a " + join_english(channels, lambda c: named[c], conj=conj) + " image"
            )

    if plural:
        return f"images with {join_english(channels, conj=conj)} channel(s)"
    else:
        return f"an image with {join_english(channels, conj=conj)} channel(s)"


def format_color_with_channels(
    channels: List[int],
    conj: Conj = "and",
    plural: bool = False,
) -> str:
    assert len(channels) > 0

    named = {1: "grayscale", 3: "RGB", 4: "RGBA"}
    if all(x in named for x in channels):
        if plural:
            return join_english(channels, lambda c: named[c], conj=conj) + " colors"
        else:
            return (
                "a " + join_english(channels, lambda c: named[c], conj=conj) + " color"
            )

    if plural:
        return f"color with {join_english(channels, conj=conj)} channel(s)"
    else:
        return f"a color with {join_english(channels, conj=conj)} channel(s)"


_CHANNEL_NUMBER_NAME = {1: "GRAY", 3: "RGB", 4: "RGBA"}


def format_channel_numbers(input_channels: int, output_channels: int) -> str:
    i = _CHANNEL_NUMBER_NAME.get(input_channels, str(input_channels))
    o = _CHANNEL_NUMBER_NAME.get(output_channels, str(output_channels))
    return f"{i}🠚{o}"

nodes/utils/format.py
vladmandic/sd-extension-chainner
from __future__ import annotations

import re


class ReplacementInterpolation:
    def __init__(self, name: str):
        self.name = name


class ReplacementString:
    """
    A parser and interpolator for chainner's string replacement patterns.

    The syntax is as follows (ANTLR 4):

        Pattern: ( LiteralChar | EscapedChar | Interpolation )* ;
        LiteralChar: ~[{] ; // any character except "{"
        EscapedChar: '{{' ;
        Interpolation: '{' InterpolationContent '}' ;
        InterpolationContent: [A-Za-z0-9]+ ;
    """

    def __init__(self, pattern: str):
        self.tokens: list[str | ReplacementInterpolation] = []
        self.names: set[str] = set()

        content_pattern = re.compile(r"\A\w+\Z")

        lastIndex = 0
        lastStr = ""
        for m in re.compile(r"(\{\{)|\{([^{}]*)\}").finditer(pattern):
            lastStr += pattern[lastIndex : m.start()]
            lastIndex = m.end()

            interpolation = m.group(2)
            if interpolation is not None:
                if interpolation == "":
                    raise ValueError(
                        "Invalid replacement pattern. {} is not a valid replacement."
                        """ Either specify a name or id number, or escape a single "{" as "{{"."""
                        f" Full pattern: {pattern}"
                    )
                if content_pattern.fullmatch(interpolation) is None:
                    raise ValueError(
                        "Invalid replacement pattern."
                        f" {{{interpolation}}} is not a valid replacement."
                        " Names and ids only allow letters and digits."
                        f" Full pattern: {pattern}"
                    )

                self.tokens.append(lastStr)
                lastStr = ""
                self.tokens.append(ReplacementInterpolation(interpolation))
                self.names.add(interpolation)
            else:
                lastStr += "{"
        lastStr += pattern[lastIndex:]
        self.tokens.append(lastStr)

    def replace(self, replacements: dict[str, str]) -> str:
        result = ""

        for token in self.tokens:
            if isinstance(token, str):
                result += token
            else:
                if token.name in replacements:
                    result += replacements[token.name]
                else:
                    raise ValueError(
                        "Unknown replacement."
                        f" There is no replacement with the name or id {token.name}."
                        f" Available replacements: {', '.join(replacements.keys()) }."
                    )

        return result

nodes/utils/replacement.py
vladmandic/sd-extension-chainner
from dataclasses import dataclass
from random import Random

_U32_MAX = 4294967296


@dataclass(frozen=True)
class Seed:
    value: int
    """
    The value of the seed. This value may be signed and generally have any range.
    """

    @staticmethod
    def from_bytes(b: bytes):
        return Seed(Random(b).randint(0, _U32_MAX - 1))

    def to_range(self, a: int, b: int) -> int:
        """
        Returns the value of the seed within the given range [a,b] both ends inclusive.

        If the current seed is not within the given range, a value within the range will be derived from the current seed.
        """
        if a <= self.value <= b:
            return self.value
        return Random(self.value).randint(a, b)

    def to_u32(self) -> int:
        """
        Returns the value of the seed as a 32bit unsigned integer.
        """
        return self.to_range(0, _U32_MAX - 1)

    def cache_key_func(self):
        return self.value

nodes/utils/seed.py
vladmandic/sd-extension-chainner
# Safe unpickler to prevent arbitrary code execution
import pickle
from types import SimpleNamespace

safe_list = {
    ("collections", "OrderedDict"),
    ("typing", "OrderedDict"),
    ("torch._utils", "_rebuild_tensor_v2"),
    ("torch", "BFloat16Storage"),
    ("torch", "FloatStorage"),
    ("torch", "HalfStorage"),
    ("torch", "IntStorage"),
    ("torch", "LongStorage"),
    ("torch", "DoubleStorage"),
}


class RestrictedUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        # Only allow required classes to load state dict
        if (module, name) not in safe_list:
            raise pickle.UnpicklingError(
                "Global '{}.{}' is forbidden".format(module, name)
            )
        return super().find_class(module, name)


RestrictedUnpickle = SimpleNamespace(
    Unpickler=RestrictedUnpickler,
    __name__="pickle",
    load=lambda *args, **kwargs: RestrictedUnpickler(*args, **kwargs).load(),
)

nodes/utils/unpickler.py
vladmandic/sd-extension-chainner
# From https://github.com/victorca25/iNNfer/blob/main/utils/utils.py
from __future__ import annotations

import math
import os
import re
from dataclasses import dataclass
from typing import List, Tuple, Union

import numpy as np
from nodes.log import logger

Size = Tuple[int, int]
"""
The width and height (in that order) of an image.
"""

NUMBERS = re.compile(r"(\d+)")

ALPHABET = [*"ABCDEFGHIJKLMNOPQRSTUVWXYZ"]


def round_half_up(number: Union[float, int]) -> int:
    """
    Python's `round` method implements round-half-to-even rounding which is very unintuitive.
    This function implements round-half-up rounding.

    Round half up is consistent with JavaScript's `Math.round`.

    https://en.wikipedia.org/wiki/Rounding#Rounding_to_the_nearest_integer
    """
    return math.floor(number + 0.5)


def get_h_w_c(image: np.ndarray) -> Tuple[int, int, int]:
    """Returns the height, width, and number of channels."""
    h, w = image.shape[:2]
    c = 1 if image.ndim == 2 else image.shape[2]
    return h, w, c


def alphanumeric_sort(value: str) -> List[Union[str, int]]:
    """Key function to sort strings containing numbers by proper
    numerical order."""

    lcase_value = value.upper()
    parts = NUMBERS.split(lcase_value)
    parts[1::2] = map(int, parts[1::2])
    return parts  # type: ignore


__SPLIT_SNAKE_CASE = re.compile(r"(\d+|_+)")
__SPLIT_PASCAL_CASE = re.compile(r"(\d+)|(?<=[a-z])(?=[A-Z])")


def split_snake_case(s: str) -> List[str]:
    """Splits a snake case identifier into its parts. E.g. `SNAKE_CASE` -> [`snake`, `case`]"""
    return [
        x.lower() for x in __SPLIT_SNAKE_CASE.split(s) if x and not x.startswith("_")
    ]


def split_pascal_case(s: str) -> List[str]:
    """Splits a snake case identifier into its parts. E.g. `SNAKE_CASE` -> [`snake`, `case`]"""
    return [
        x.lower() for x in __SPLIT_PASCAL_CASE.split(s) if x and not x.startswith("_")
    ]


def join_pascal_case(words: List[str]) -> str:
    return "".join([x.capitalize() for x in words])


__ABBREVIATIONS = {"rgb", "rgba"}


def smart_capitalize(word: str) -> str:
    if word in __ABBREVIATIONS:
        return word.upper()
    return word.capitalize()


def join_space_case(words: List[str]) -> str:
    return " ".join([smart_capitalize(x) for x in words])


def split_file_path(path: str) -> Tuple[str, str, str]:
    """
    Returns the base directory, file name, and extension of the given file path.
    """
    base, ext = os.path.splitext(path)
    dirname, basename = os.path.split(base)
    return dirname, basename, ext


def walk_error_handler(exception_instance):
    logger.warning(
        f"Exception occurred during walk: {exception_instance} Continuing..."
    )


def list_all_files_sorted(
    directory: str, ext_filter: Union[List[str], None] = None
) -> List[str]:
    just_files: List[str] = []
    for root, dirs, files in os.walk(
        directory, topdown=True, onerror=walk_error_handler
    ):
        dirs.sort(key=alphanumeric_sort)
        for name in sorted(files, key=alphanumeric_sort):
            filepath = os.path.join(root, name)
            _base, ext = os.path.splitext(filepath)
            if ext_filter is None or ext.lower() in ext_filter:
                just_files.append(filepath)
    return just_files


@dataclass(frozen=True)
class Padding:
    top: int
    right: int
    bottom: int
    left: int

    @staticmethod
    def all(value: int) -> "Padding":
        return Padding(value, value, value, value)

    @staticmethod
    def to(value: Padding | int) -> Padding:
        if isinstance(value, int):
            return Padding.all(value)
        return value

    @property
    def horizontal(self) -> int:
        return self.left + self.right

    @property
    def vertical(self) -> int:
        return self.top + self.bottom

    @property
    def empty(self) -> bool:
        return self.top == 0 and self.right == 0 and self.bottom == 0 and self.left == 0

    def scale(self, factor: int) -> Padding:
        return Padding(
            self.top * factor,
            self.right * factor,
            self.bottom * factor,
            self.left * factor,
        )

    def min(self, other: Padding | int) -> Padding:
        other = Padding.to(other)
        return Padding(
            min(self.top, other.top),
            min(self.right, other.right),
            min(self.bottom, other.bottom),
            min(self.left, other.left),
        )

    def remove_from(self, image: np.ndarray) -> np.ndarray:
        h, w, _ = get_h_w_c(image)

        return image[
            self.top : (h - self.bottom),
            self.left : (w - self.right),
            ...,
        ]


@dataclass(frozen=True)
class Region:
    x: int
    y: int
    width: int
    height: int

    @property
    def size(self) -> Size:
        return self.width, self.height

    def scale(self, factor: int) -> Region:
        return Region(
            self.x * factor,
            self.y * factor,
            self.width * factor,
            self.height * factor,
        )

    def intersect(self, other: Region) -> Region:
        x = max(self.x, other.x)
        y = max(self.y, other.y)
        width = min(self.x + self.width, other.x + other.width) - x
        height = min(self.y + self.height, other.y + other.height) - y
        return Region(x, y, width, height)

    def add_padding(self, pad: Padding) -> Region:
        return Region(
            x=self.x - pad.left,
            y=self.y - pad.top,
            width=self.width + pad.horizontal,
            height=self.height + pad.vertical,
        )

    def remove_padding(self, pad: Padding) -> Region:
        return self.add_padding(pad.scale(-1))

    def child_padding(self, child: Region) -> Padding:
        """
        Returns the padding `p` such that `child.add_padding(p) == self`.
        """
        left = child.x - self.x
        top = child.y - self.y
        right = self.width - child.width - left
        bottom = self.height - child.height - top
        return Padding(top, right, bottom, left)

    def read_from(self, image: np.ndarray) -> np.ndarray:
        h, w, _ = get_h_w_c(image)
        if (w, h) == self.size:
            return image

        return image[
            self.y : (self.y + self.height),
            self.x : (self.x + self.width),
            ...,
        ]

    def write_into(self, lhs: np.ndarray, rhs: np.ndarray):
        h, w, c = get_h_w_c(rhs)
        assert (w, h) == self.size
        assert c == get_h_w_c(lhs)[2]

        if c == 1:
            if lhs.ndim == 2 and rhs.ndim == 3:
                rhs = rhs[:, :, 0]
            if lhs.ndim == 3 and rhs.ndim == 2:
                rhs = np.expand_dims(rhs, axis=2)

        lhs[
            self.y : (self.y + self.height),
            self.x : (self.x + self.width),
            ...,
        ] = rhs

nodes/utils/utils.py
vladmandic/sd-extension-chainner
import os
import PIL.Image
import numpy as np
import torch
from nodes.impl.upscale.tiler import MaxTileSize, NoTiling, Tiler
from nodes.impl.pytorch.auto_split import pytorch_auto_split
from nodes.load_model import load_model
from nodes.impl.pytorch.types import PyTorchSRModel
from modules import devices, script_callbacks # pylint: disable=wrong-import-order
from modules.shared import opts, log, paths, OptionInfo # pylint: disable=wrong-import-order
from modules.upscaler import Upscaler, UpscalerData # pylint: disable=wrong-import-order


chainner_models = [
    ["4x DAT", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/DAT-4x.pth"],
    ["2x HAT", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/HAT-2x.pth"],
    ["3x HAT", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/HAT-3x.pth"],
    ["4x HAT", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/HAT-4x.pth"],
    ["2x HAT-Large", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/HAT-L-2x.pth"],
    ["3x HAT-Large", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/HAT-L-3x.pth"],
    ["4x HAT-Large", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/HAT-L-4x.pth"],
    ["4x RRDBNet", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/RRDBNet-4x.pth"],
    ["4x RealHAT-GAN", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/RealHAT-GAN-4x.pth"],
    ["4x RealHAT-Sharper", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/RealHAT-Sharper-4x.pth"],
    ["4x SPSRNet", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/SPSRNet-4x.pth"],
    ["4x SRFormer-Light", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/SRFormer-Light-4x.pth"],
    ["4x SRFormer-Nomos", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/SRFormer-Nomos-4x.pth"],
    ["2x SwiftSR", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/SwiftSR-2x.pth"],
    ["4x SwiftSR", "https://huggingface.co/vladmandic/sdnext-upscalers/resolve/main/SwiftSR-4x.pth"],
]


def friendly_fullname(file: str):
    from urllib.parse import urlparse
    if "http" in file:
        file = urlparse(file).path
    file = os.path.basename(file)
    return file


class UpscalerChaiNNer(Upscaler):
    def __init__(self, dirname): # pylint: disable=unused-argument
        self.name = "chaiNNer"
        self.user_path = opts.data.get('chainner_models_path', os.path.join(paths.models_path, 'chaiNNer'))
        super().__init__()
        self.models = {}
        self.fp16 = False
        self.scalers = self.find_scalers()

    def find_scalers(self):
        loaded = []
        scalers = []
        for model in chainner_models:
            local_name = os.path.join(self.user_path, friendly_fullname(model[1]))
            model_path = local_name if os.path.exists(local_name) else model[1]
            scaler = UpscalerData(name=f'{self.name} {model[0]}', path=model_path, upscaler=self)
            scalers.append(scaler)
            loaded.append(model_path)
        for fn in os.listdir(self.user_path): # from folder
            if not fn.endswith('.pth'):
                continue
            file_name = os.path.join(self.user_path, fn)
            if file_name not in loaded:
                model_name = os.path.splitext(fn)[0]
                scaler = UpscalerData(name=f'{self.name} {model_name}', path=file_name, upscaler=self)
                scaler.custom = True
                scalers.append(scaler)
                loaded.append(file_name)
        return scalers

    def load_model(self, path: str):
        info = self.find_model(path)
        if info is None:
            return
        if self.models.get(info.local_data_path, None) is not None:
            log.debug(f"Upscaler cached: type={self.name} model={info.local_data_path}")
            model=self.models[info.local_data_path]
        else:
            model: PyTorchSRModel = load_model(info.local_data_path, device=devices.device, fp16=self.fp16)
            log.info(f"Upscaler loaded: type={self.name} model='{info.local_data_path}'")
            self.models[info.local_data_path] = model
        return model

    def parse_tile_size_input(self, tile_size: int) -> Tiler:
        if tile_size == 0:
            return MaxTileSize(tile_size)
        elif tile_size == -1:
            return NoTiling()
        elif tile_size == -2:
            return MaxTileSize()
        elif tile_size < 0:
            raise ValueError(f"ChaiNNer invalid tile size: {tile_size}")
        return MaxTileSize(tile_size)

    def do_upscale(self, img: PIL.Image.Image, selected_model):
        devices.torch_gc()
        model = self.load_model(selected_model)
        if model is None:
            return img
        np_img = np.array(img)
        tile_size = opts.data.get('upscaler_tile_size', 192)
        with torch.no_grad():
            upscaled = pytorch_auto_split(img=np_img, model=model, device=devices.device, use_fp16=self.fp16, tiler=self.parse_tile_size_input(tile_size))
            img = PIL.Image.fromarray(np.uint8(250.0 * upscaled)) # TODO full range causes some color clipping?
        devices.torch_gc()
        if opts.data.get('upscaler_unload', False) and selected_model in self.models:
            del self.models[selected_model]
            log.debug(f"Upscaler unloaded: type={self.name} model={selected_model}")
            devices.torch_gc(force=True)
        return img


def on_ui_settings():
    opts.add_option("chainner_models_path", OptionInfo(os.path.join(paths.models_path, 'chaiNNer'), "Folder with chaiNNer models", folder=True, section=('system-paths', "System Paths")))


script_callbacks.on_ui_settings(on_ui_settings)

scripts/chainner_model.py
vladmandic/sd-extension-chainner