chore: import upstream snapshot with attribution

invokeai/backend/image_util/__init__.py

@@ -0,0 +1,12 @@
+"""
+Initialization file for invokeai.backend.image_util methods.
+"""
+
+from invokeai.backend.image_util.infill_methods.patchmatch import PatchMatch  # noqa: F401
+from invokeai.backend.image_util.pngwriter import (  # noqa: F401
+    PngWriter,
+    PromptFormatter,
+    retrieve_metadata,
+    write_metadata,
+)
+from invokeai.backend.image_util.util import InitImageResizer, make_grid  # noqa: F401

invokeai/backend/image_util/basicsr/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2018-2022 BasicSR Authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

invokeai/backend/image_util/basicsr/__init__.py

@@ -0,0 +1,18 @@
+"""
+Adapted from https://github.com/XPixelGroup/BasicSR
+License: Apache-2.0
+
+As of Feb 2024, `basicsr` appears to be unmaintained. It imports a function from `torchvision` that is removed in
+`torchvision` 0.17. Here is the deprecation warning:
+
+    UserWarning: The torchvision.transforms.functional_tensor module is deprecated in 0.15 and will be **removed in
+    0.17**. Please don't rely on it. You probably just need to use APIs in torchvision.transforms.functional or in
+    torchvision.transforms.v2.functional.
+
+As a result, a dependency on `basicsr` means we cannot keep our `torchvision` dependency up to date.
+
+Because we only rely on a single class `RRDBNet` from `basicsr`, we've copied the relevant code here and removed the
+dependency on `basicsr`.
+
+The code is almost unchanged, only a few type annotations have been added. The license is also copied.
+"""

invokeai/backend/image_util/basicsr/arch_util.py

@@ -0,0 +1,75 @@
+from typing import Type
+
+import torch
+from torch import nn as nn
+from torch.nn import init as init
+from torch.nn.modules.batchnorm import _BatchNorm
+
+
+@torch.no_grad()
+def default_init_weights(
+    module_list: list[nn.Module] | nn.Module, scale: float = 1, bias_fill: float = 0, **kwargs
+) -> None:
+    """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)
+
+
+def make_layer(basic_block: Type[nn.Module], num_basic_block: int, **kwarg) -> nn.Sequential:
+    """Make layers by stacking the same blocks.
+
+    Args:
+        basic_block (Type[nn.Module]): nn.Module class for basic block.
+        num_basic_block (int): number of blocks.
+
+    Returns:
+        nn.Sequential: Stacked blocks in nn.Sequential.
+    """
+    layers = []
+    for _ in range(num_basic_block):
+        layers.append(basic_block(**kwarg))
+    return nn.Sequential(*layers)
+
+
+# TODO: may write a cpp file
+def pixel_unshuffle(x: torch.Tensor, scale: int) -> torch.Tensor:
+    """Pixel unshuffle.
+
+    Args:
+        x (Tensor): Input feature with shape (b, c, hh, hw).
+        scale (int): Downsample ratio.
+
+    Returns:
+        Tensor: the pixel unshuffled feature.
+    """
+    b, c, hh, hw = x.size()
+    out_channel = c * (scale**2)
+    assert hh % scale == 0 and hw % scale == 0
+    h = hh // scale
+    w = hw // scale
+    x_view = x.view(b, c, h, scale, w, scale)
+    return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)

invokeai/backend/image_util/basicsr/rrdbnet_arch.py

@@ -0,0 +1,125 @@
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+from invokeai.backend.image_util.basicsr.arch_util import default_init_weights, make_layer, pixel_unshuffle
+
+
+class ResidualDenseBlock(nn.Module):
+    """Residual Dense Block.
+
+    Used in RRDB block in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat: int = 64, num_grow_ch: int = 32) -> None:
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
+        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        # Empirically, we use 0.2 to scale the residual for better performance
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    """Residual in Residual Dense Block.
+
+    Used in RRDB-Net in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat: int, num_grow_ch: int = 32) -> None:
+        super(RRDB, self).__init__()
+        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        # Empirically, we use 0.2 to scale the residual for better performance
+        return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    """Networks consisting of Residual in Residual Dense Block, which is used
+    in ESRGAN.
+
+    ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.
+
+    We extend ESRGAN for scale x2 and scale x1.
+    Note: This is one option for scale 1, scale 2 in RRDBNet.
+    We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size
+    and enlarge the channel size before feeding inputs into the main ESRGAN architecture.
+
+    Args:
+        num_in_ch (int): Channel number of inputs.
+        num_out_ch (int): Channel number of outputs.
+        num_feat (int): Channel number of intermediate features.
+            Default: 64
+        num_block (int): Block number in the trunk network. Defaults: 23
+        num_grow_ch (int): Channels for each growth. Default: 32.
+    """
+
+    def __init__(
+        self,
+        num_in_ch: int,
+        num_out_ch: int,
+        scale: int = 4,
+        num_feat: int = 64,
+        num_block: int = 23,
+        num_grow_ch: int = 32,
+    ) -> None:
+        super(RRDBNet, self).__init__()
+        self.scale = scale
+        if scale == 2:
+            num_in_ch = num_in_ch * 4
+        elif scale == 1:
+            num_in_ch = num_in_ch * 16
+        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
+        self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
+        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        # upsample
+        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)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.scale == 2:
+            feat = pixel_unshuffle(x, scale=2)
+        elif self.scale == 1:
+            feat = pixel_unshuffle(x, scale=4)
+        else:
+            feat = x
+        feat = self.conv_first(feat)
+        body_feat = self.conv_body(self.body(feat))
+        feat = feat + body_feat
+        # upsample
+        feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode="nearest")))
+        feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode="nearest")))
+        out = self.conv_last(self.lrelu(self.conv_hr(feat)))
+        return out

invokeai/backend/image_util/canny.py

@@ -0,0 +1,41 @@
+import cv2
+from PIL import Image
+
+from invokeai.backend.image_util.util import (
+    cv2_to_pil,
+    normalize_image_channel_count,
+    pil_to_cv2,
+    resize_image_to_resolution,
+)
+
+
+def get_canny_edges(
+    image: Image.Image, low_threshold: int, high_threshold: int, detect_resolution: int, image_resolution: int
+) -> Image.Image:
+    """Returns the edges of an image using the Canny edge detection algorithm.
+
+    Adapted from https://github.com/huggingface/controlnet_aux (Apache-2.0 license).
+
+    Args:
+        image: The input image.
+        low_threshold: The lower threshold for the hysteresis procedure.
+        high_threshold: The upper threshold for the hysteresis procedure.
+        input_resolution: The resolution of the input image. The image will be resized to this resolution before edge detection.
+        output_resolution: The resolution of the output image. The edges will be resized to this resolution before returning.
+
+    Returns:
+        The Canny edges of the input image.
+    """
+
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    np_image = pil_to_cv2(image)
+    np_image = normalize_image_channel_count(np_image)
+    np_image = resize_image_to_resolution(np_image, detect_resolution)
+
+    edge_map = cv2.Canny(np_image, low_threshold, high_threshold)
+    edge_map = normalize_image_channel_count(edge_map)
+    edge_map = resize_image_to_resolution(edge_map, image_resolution)
+
+    return cv2_to_pil(edge_map)

invokeai/backend/image_util/composition.py

@@ -0,0 +1,122 @@
+# TODO: Improve blend modes
+# TODO: Add nodes like Hue Adjust for Saturation/Contrast/etc... ?
+# TODO: Continue implementing more blend modes/color spaces(?)
+# TODO: Custom ICC profiles with PIL.ImageCms?
+# TODO: Blend multiple layers all crammed into a tensor(?) or list
+
+# Copyright (c) 2023 Darren Ringer <dwringer@gmail.com>
+# Parts based on Oklab: Copyright (c) 2021 Bj�rn Ottosson <https://bottosson.github.io/>
+# HSL code based on CPython: Copyright (c) 2001-2023 Python Software Foundation; All Rights Reserved
+from math import pi as PI
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+from invokeai.backend.image_util.color_conversion import (
+    gamut_clip_tensor,
+)
+from invokeai.backend.image_util.color_conversion import (
+    srgb_from_linear_srgb as shared_srgb_from_linear_srgb,
+)
+from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
+
+MAX_FLOAT = torch.finfo(torch.tensor(1.0).dtype).max
+
+# CIE Lab to Uniform Perceptual Lab profile is copyright © 2003 Bruce Justin Lindbloom. All rights reserved. <http://www.brucelindbloom.com>
+CIELAB_TO_UPLAB_ICC_PATH = Path(__file__).parent / "assets" / "CIELab_to_UPLab.icc"
+
+
+def equivalent_achromatic_lightness(lch_tensor: torch.Tensor):
+    """Calculate Equivalent Achromatic Lightness accounting for Helmholtz-Kohlrausch effect"""
+    # As described by High, Green, and Nussbaum (2023): https://doi.org/10.1002/col.22839
+
+    k = [0.1644, 0.0603, 0.1307, 0.0060]
+
+    h_minus_90 = torch.sub(lch_tensor[2, :, :], PI / 2.0)
+    h_minus_90 = torch.sub(torch.remainder(torch.add(h_minus_90, 3 * PI), 2 * PI), PI)
+
+    f_by = torch.add(k[0] * torch.abs(torch.sin(torch.div(h_minus_90, 2.0))), k[1])
+    f_r_0 = torch.add(k[2] * torch.abs(torch.cos(lch_tensor[2, :, :])), k[3])
+
+    f_r = torch.zeros(lch_tensor[0, :, :].shape)
+    mask_hi = torch.ge(lch_tensor[2, :, :], -1 * (PI / 2.0))
+    mask_lo = torch.le(lch_tensor[2, :, :], PI / 2.0)
+    mask = torch.logical_and(mask_hi, mask_lo)
+    f_r[mask] = f_r_0[mask]
+
+    l_max = torch.ones(lch_tensor[0, :, :].shape)
+    l_min = torch.zeros(lch_tensor[0, :, :].shape)
+    l_adjustment = torch.tensordot(torch.add(f_by, f_r), lch_tensor[1, :, :], dims=([0, 1], [0, 1]))
+    l_max = torch.add(l_max, l_adjustment)
+    l_min = torch.add(l_min, l_adjustment)
+    l_eal_tensor = torch.add(lch_tensor[0, :, :], l_adjustment)
+
+    l_eal_tensor = torch.add(
+        lch_tensor[0, :, :], torch.tensordot(torch.add(f_by, f_r), lch_tensor[1, :, :], dims=([0, 1], [0, 1]))
+    )
+    l_eal_tensor = torch.div(torch.sub(l_eal_tensor, l_min.min()), l_max.max() - l_min.min())
+
+    return l_eal_tensor
+
+
+def srgb_from_linear_srgb(linear_srgb_tensor: torch.Tensor, alpha: float = 0.0, steps: int = 1):
+    """Get gamma-corrected sRGB from a linear-light sRGB image tensor"""
+
+    if 0.0 < alpha:
+        linear_srgb_tensor = gamut_clip_tensor(linear_srgb_tensor, alpha=alpha, steps=steps)
+    return shared_srgb_from_linear_srgb(linear_srgb_tensor)
+
+
+def remove_nans(tensor: torch.Tensor, replace_with: float = MAX_FLOAT):
+    return torch.where(torch.isnan(tensor), replace_with, tensor)
+
+
+def tensor_from_pil_image(img: Image.Image, normalize: bool = False):
+    return image_resized_to_grid_as_tensor(img, normalize=normalize, multiple_of=1)
+
+
+# PSF LICENSE AGREEMENT FOR PYTHON 3.11.5
+
+# 1. This LICENSE AGREEMENT is between the Python Software Foundation ("PSF"), and
+#    the Individual or Organization ("Licensee") accessing and otherwise using Python
+#    3.11.5 software in source or binary form and its associated documentation.
+
+# 2. Subject to the terms and conditions of this License Agreement, PSF hereby
+#    grants Licensee a nonexclusive, royalty-free, world-wide license to reproduce,
+#    analyze, test, perform and/or display publicly, prepare derivative works,
+#    distribute, and otherwise use Python 3.11.5 alone or in any derivative
+#    version, provided, however, that PSF's License Agreement and PSF's notice of
+#    copyright, i.e., "Copyright (c) 2001-2023 Python Software Foundation; All Rights
+#    Reserved" are retained in Python 3.11.5 alone or in any derivative version
+#    prepared by Licensee.
+
+# 3. In the event Licensee prepares a derivative work that is based on or
+#    incorporates Python 3.11.5 or any part thereof, and wants to make the
+#    derivative work available to others as provided herein, then Licensee hereby
+#    agrees to include in any such work a brief summary of the changes made to Python
+#    3.11.5.
+
+# 4. PSF is making Python 3.11.5 available to Licensee on an "AS IS" basis.
+#    PSF MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR IMPLIED.  BY WAY OF
+#    EXAMPLE, BUT NOT LIMITATION, PSF MAKES NO AND DISCLAIMS ANY REPRESENTATION OR
+#    WARRANTY OF MERCHANTABILITY OR FITNESS FOR ANY PARTICULAR PURPOSE OR THAT THE
+#    USE OF PYTHON 3.11.5 WILL NOT INFRINGE ANY THIRD PARTY RIGHTS.
+
+# 5. PSF SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON 3.11.5
+#    FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS A RESULT OF
+#    MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON 3.11.5, OR ANY DERIVATIVE
+#    THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF.
+
+# 6. This License Agreement will automatically terminate upon a material breach of
+#    its terms and conditions.
+
+# 7. Nothing in this License Agreement shall be deemed to create any relationship
+#    of agency, partnership, or joint venture between PSF and Licensee.  This License
+#    Agreement does not grant permission to use PSF trademarks or trade name in a
+#    trademark sense to endorse or promote products or services of Licensee, or any
+#    third party.
+
+# 8. By copying, installing or otherwise using Python 3.11.5, Licensee agrees
+#    to be bound by the terms and conditions of this License Agreement.
+######################################################################################/

invokeai/backend/image_util/content_shuffle.py

@@ -0,0 +1,40 @@
+# Adapted from https://github.com/huggingface/controlnet_aux
+
+import cv2
+import numpy as np
+from PIL import Image
+
+from invokeai.backend.image_util.util import np_to_pil, pil_to_np
+
+
+def make_noise_disk(H, W, C, F):
+    noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))
+    noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)
+    noise = noise[F : F + H, F : F + W]
+    noise -= np.min(noise)
+    noise /= np.max(noise)
+    if C == 1:
+        noise = noise[:, :, None]
+    return noise
+
+
+def content_shuffle(input_image: Image.Image, scale_factor: int | None = None) -> Image.Image:
+    """Shuffles the content of an image using a disk noise pattern, similar to a 'liquify' effect."""
+
+    np_img = pil_to_np(input_image)
+
+    height, width, _channels = np_img.shape
+
+    if scale_factor is None:
+        scale_factor = 256
+
+    x = make_noise_disk(height, width, 1, scale_factor) * float(width - 1)
+    y = make_noise_disk(height, width, 1, scale_factor) * float(height - 1)
+
+    flow = np.concatenate([x, y], axis=2).astype(np.float32)
+
+    shuffled_img = cv2.remap(np_img, flow, None, cv2.INTER_LINEAR)
+
+    output_img = np_to_pil(shuffled_img)
+
+    return output_img

invokeai/backend/image_util/controlnet_processor.py

@@ -0,0 +1,188 @@
+"""Utilities for processing images with ControlNet processors."""
+
+from datetime import datetime
+from typing import Any, Optional
+
+from invokeai.app.invocations.fields import ImageField
+from invokeai.app.services.invoker import InvocationServices
+from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem
+from invokeai.app.services.shared.graph import Graph, GraphExecutionState
+from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context
+
+
+def _get_processor_invocation_class(processor_type: str):
+    """Get the invocation class for a processor type."""
+    # Import processor invocation classes on demand
+    processor_class_map = {
+        "canny_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.canny", fromlist=["CannyEdgeDetectionInvocation"]
+            ).CannyEdgeDetectionInvocation
+        ),
+        "hed_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.hed", fromlist=["HEDEdgeDetectionInvocation"]
+            ).HEDEdgeDetectionInvocation
+        ),
+        "mlsd_image_processor": lambda: (
+            __import__("invokeai.app.invocations.mlsd", fromlist=["MLSDDetectionInvocation"]).MLSDDetectionInvocation
+        ),
+        "depth_anything_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.depth_anything", fromlist=["DepthAnythingDepthEstimationInvocation"]
+            ).DepthAnythingDepthEstimationInvocation
+        ),
+        "normalbae_image_processor": lambda: (
+            __import__("invokeai.app.invocations.normal_bae", fromlist=["NormalMapInvocation"]).NormalMapInvocation
+        ),
+        "pidi_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.pidi", fromlist=["PiDiNetEdgeDetectionInvocation"]
+            ).PiDiNetEdgeDetectionInvocation
+        ),
+        "lineart_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.lineart", fromlist=["LineartEdgeDetectionInvocation"]
+            ).LineartEdgeDetectionInvocation
+        ),
+        "lineart_anime_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.lineart_anime", fromlist=["LineartAnimeEdgeDetectionInvocation"]
+            ).LineartAnimeEdgeDetectionInvocation
+        ),
+        "content_shuffle_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.content_shuffle", fromlist=["ContentShuffleInvocation"]
+            ).ContentShuffleInvocation
+        ),
+        "dw_openpose_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.dw_openpose", fromlist=["DWOpenposeDetectionInvocation"]
+            ).DWOpenposeDetectionInvocation
+        ),
+        "mediapipe_face_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.mediapipe_face", fromlist=["MediaPipeFaceDetectionInvocation"]
+            ).MediaPipeFaceDetectionInvocation
+        ),
+        # Note: zoe_depth_image_processor doesn't have a processor invocation implementation
+        "color_map_image_processor": lambda: (
+            __import__("invokeai.app.invocations.color_map", fromlist=["ColorMapInvocation"]).ColorMapInvocation
+        ),
+    }
+
+    if processor_type in processor_class_map:
+        return processor_class_map[processor_type]()
+    return None
+
+
+# Map processor type names to their default parameters
+PROCESSOR_DEFAULT_PARAMS = {
+    "canny_image_processor": {"low_threshold": 100, "high_threshold": 200},
+    "hed_image_processor": {"scribble": False},
+    "mlsd_image_processor": {"detect_resolution": 512, "thr_v": 0.1, "thr_d": 0.1},
+    "depth_anything_image_processor": {"model_size": "small"},
+    "normalbae_image_processor": {"detect_resolution": 512},
+    "pidi_image_processor": {"detect_resolution": 512, "safe": False},
+    "lineart_image_processor": {"detect_resolution": 512, "coarse": False},
+    "lineart_anime_image_processor": {"detect_resolution": 512},
+    "content_shuffle": {},
+    "dw_openpose_image_processor": {"draw_body": True, "draw_face": True, "draw_hands": True},
+    "mediapipe_face_processor": {"max_faces": 1, "min_confidence": 0.5},
+    "zoe_depth_image_processor": {},
+    "color_map_image_processor": {"color_map_tile_size": 64},
+}
+
+
+def process_controlnet_image(image_name: str, model_key: str, services: InvocationServices) -> Optional[dict[str, Any]]:
+    """
+    Process a controlnet image using the appropriate processor based on the model's default settings.
+
+    Args:
+        image_name: The filename of the image to process
+        model_key: The model key to look up default processor settings
+        services: The invocation services providing access to models and images
+
+    Returns:
+        A dictionary with the processed image data (image_name, width, height) or None if processing fails
+    """
+    logger = services.logger
+
+    try:
+        # Get model config to find default processor
+        model_record = services.model_manager.store.get_model(model_key)
+        if not model_record or not model_record.default_settings:
+            logger.info(f"No default processor settings found for model {model_key}")
+            return None
+
+        preprocessor = model_record.default_settings.preprocessor
+        if not preprocessor:
+            logger.info(f"No preprocessor configured for model {model_key}")
+            return None
+
+        # Get the invocation class for this processor
+        invocation_class = _get_processor_invocation_class(preprocessor)
+        if not invocation_class:
+            logger.info(f"No processor mapping found for preprocessor '{preprocessor}'")
+            return None
+
+        # Get default parameters for this processor
+        default_params = PROCESSOR_DEFAULT_PARAMS.get(preprocessor, {})
+        logger.info(f"Processing image {image_name} with processor {preprocessor}")
+
+        # Create a minimal context to run the invocation
+        # We need a fake queue item and session for the context
+        fake_session = GraphExecutionState(graph=Graph())
+        now = datetime.now()
+
+        # Create invocation instance first so we have its ID
+        invocation_params = {"image": ImageField(image_name=image_name), **default_params}
+        invocation = invocation_class(**invocation_params)
+
+        # Add the invocation ID to the session's prepared_source_mapping
+        # This is required for the invocation context to emit progress events
+        fake_session.prepared_source_mapping[invocation.id] = invocation.id
+
+        fake_queue_item = SessionQueueItem(
+            item_id=0,
+            session_id=fake_session.id,
+            queue_id="default",
+            batch_id="recall_processor",
+            field_values=None,
+            session=fake_session,
+            status="in_progress",
+            created_at=now,
+            updated_at=now,
+            started_at=now,
+            completed_at=None,
+        )
+
+        context_data = InvocationContextData(
+            invocation=invocation,
+            source_invocation_id=invocation.id,
+            queue_item=fake_queue_item,
+        )
+
+        context = build_invocation_context(
+            data=context_data,
+            services=services,
+            is_canceled=lambda: False,
+        )
+
+        # Invoke the processor
+        output = invocation.invoke(context)
+
+        # Get the processed image DTO
+        processed_image_dto = services.images.get_dto(output.image.image_name)
+
+        logger.info(f"Successfully processed image {image_name} -> {processed_image_dto.image_name}")
+
+        return {
+            "image_name": processed_image_dto.image_name,
+            "width": processed_image_dto.width,
+            "height": processed_image_dto.height,
+        }
+
+    except Exception as e:
+        logger.error(f"Error processing controlnet image {image_name}: {e}", exc_info=True)
+        return None

invokeai/backend/image_util/depth_anything/depth_anything_pipeline.py

@@ -0,0 +1,41 @@
+import pathlib
+from typing import Optional
+
+import torch
+from PIL import Image
+from transformers import pipeline
+from transformers.pipelines import DepthEstimationPipeline
+
+from invokeai.backend.raw_model import RawModel
+
+
+class DepthAnythingPipeline(RawModel):
+    """Custom wrapper for the Depth Estimation pipeline from transformers adding compatibility
+    for Invoke's Model Management System"""
+
+    def __init__(self, pipeline: DepthEstimationPipeline) -> None:
+        self._pipeline = pipeline
+
+    def generate_depth(self, image: Image.Image) -> Image.Image:
+        depth_map = self._pipeline(image)["depth"]
+        assert isinstance(depth_map, Image.Image)
+        return depth_map
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        if device is not None and device.type not in {"cpu", "cuda"}:
+            device = None
+        self._pipeline.model.to(device=device, dtype=dtype)
+        self._pipeline.device = self._pipeline.model.device
+
+    def calc_size(self) -> int:
+        from invokeai.backend.model_manager.load.model_util import calc_module_size
+
+        return calc_module_size(self._pipeline.model)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path):
+        """Load the model from the given path and return a DepthAnythingPipeline instance."""
+
+        depth_anything_pipeline = pipeline(model=str(model_path), task="depth-estimation", local_files_only=True)
+        assert isinstance(depth_anything_pipeline, DepthEstimationPipeline)
+        return cls(depth_anything_pipeline)

invokeai/backend/image_util/dw_openpose/__init__.py

@@ -0,0 +1,152 @@
+from pathlib import Path
+from typing import Dict
+
+import huggingface_hub
+import numpy as np
+import onnxruntime as ort
+import torch
+from PIL import Image
+
+from invokeai.backend.image_util.dw_openpose.onnxdet import inference_detector
+from invokeai.backend.image_util.dw_openpose.onnxpose import inference_pose
+from invokeai.backend.image_util.dw_openpose.utils import NDArrayInt, draw_bodypose, draw_facepose, draw_handpose
+from invokeai.backend.image_util.util import np_to_pil
+from invokeai.backend.util.devices import TorchDevice
+
+
+class DWOpenposeDetector:
+    """
+    Code from the original implementation of the DW Openpose Detector.
+    Credits: https://github.com/IDEA-Research/DWPose
+    """
+
+    hf_repo_id = "yzd-v/DWPose"
+    hf_filename_onnx_det = "yolox_l.onnx"
+    hf_filename_onnx_pose = "dw-ll_ucoco_384.onnx"
+
+    @classmethod
+    def get_model_url_det(cls) -> str:
+        """Returns the URL for the detection model."""
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename_onnx_det)
+
+    @classmethod
+    def get_model_url_pose(cls) -> str:
+        """Returns the URL for the pose model."""
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename_onnx_pose)
+
+    @staticmethod
+    def create_onnx_inference_session(model_path: Path) -> ort.InferenceSession:
+        """Creates an ONNX Inference Session for the given model path, using the appropriate execution provider based on
+        the device type."""
+
+        device = TorchDevice.choose_torch_device()
+        providers = ["CUDAExecutionProvider"] if device.type == "cuda" else ["CPUExecutionProvider"]
+        return ort.InferenceSession(path_or_bytes=model_path, providers=providers)
+
+    def __init__(self, session_det: ort.InferenceSession, session_pose: ort.InferenceSession):
+        self.session_det = session_det
+        self.session_pose = session_pose
+
+    def pose_estimation(self, np_image: np.ndarray):
+        """Does the pose estimation on the given image and returns the keypoints and scores."""
+
+        det_result = inference_detector(self.session_det, np_image)
+        keypoints, scores = inference_pose(self.session_pose, det_result, np_image)
+
+        keypoints_info = np.concatenate((keypoints, scores[..., None]), axis=-1)
+        # compute neck joint
+        neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
+        # neck score when visualizing pred
+        neck[:, 2:4] = np.logical_and(keypoints_info[:, 5, 2:4] > 0.3, keypoints_info[:, 6, 2:4] > 0.3).astype(int)
+        new_keypoints_info = np.insert(keypoints_info, 17, neck, axis=1)
+        mmpose_idx = [17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3]
+        openpose_idx = [1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17]
+        new_keypoints_info[:, openpose_idx] = new_keypoints_info[:, mmpose_idx]
+        keypoints_info = new_keypoints_info
+
+        keypoints, scores = keypoints_info[..., :2], keypoints_info[..., 2]
+
+        return keypoints, scores
+
+    def run(
+        self,
+        image: Image.Image,
+        draw_face: bool = False,
+        draw_body: bool = True,
+        draw_hands: bool = False,
+    ) -> Image.Image:
+        """Detects the pose in the given image and returns an solid black image with pose drawn on top, suitable for
+        use with a ControlNet."""
+
+        np_image = np.array(image)
+        H, W, C = np_image.shape
+
+        with torch.no_grad():
+            candidate, subset = self.pose_estimation(np_image)
+            nums, keys, locs = candidate.shape
+            candidate[..., 0] /= float(W)
+            candidate[..., 1] /= float(H)
+            body = candidate[:, :18].copy()
+            body = body.reshape(nums * 18, locs)
+            score = subset[:, :18]
+            for i in range(len(score)):
+                for j in range(len(score[i])):
+                    if score[i][j] > 0.3:
+                        score[i][j] = int(18 * i + j)
+                    else:
+                        score[i][j] = -1
+
+            un_visible = subset < 0.3
+            candidate[un_visible] = -1
+
+            # foot = candidate[:, 18:24]
+
+            faces = candidate[:, 24:92]
+
+            hands = candidate[:, 92:113]
+            hands = np.vstack([hands, candidate[:, 113:]])
+
+            bodies = {"candidate": body, "subset": score}
+            pose = {"bodies": bodies, "hands": hands, "faces": faces}
+
+            return DWOpenposeDetector.draw_pose(
+                pose, H, W, draw_face=draw_face, draw_hands=draw_hands, draw_body=draw_body
+            )
+
+    @staticmethod
+    def draw_pose(
+        pose: Dict[str, NDArrayInt | Dict[str, NDArrayInt]],
+        H: int,
+        W: int,
+        draw_face: bool = True,
+        draw_body: bool = True,
+        draw_hands: bool = True,
+    ) -> Image.Image:
+        """Draws the pose on a black image and returns it as a PIL Image."""
+
+        bodies = pose["bodies"]
+        faces = pose["faces"]
+        hands = pose["hands"]
+
+        assert isinstance(bodies, dict)
+        candidate = bodies["candidate"]
+
+        assert isinstance(bodies, dict)
+        subset = bodies["subset"]
+
+        canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
+
+        if draw_body:
+            canvas = draw_bodypose(canvas, candidate, subset)
+
+        if draw_hands:
+            assert isinstance(hands, np.ndarray)
+            canvas = draw_handpose(canvas, hands)
+
+        if draw_face:
+            assert isinstance(hands, np.ndarray)
+            canvas = draw_facepose(canvas, faces)  # type: ignore
+
+        dwpose_image = np_to_pil(canvas)
+
+        return dwpose_image

invokeai/backend/image_util/dw_openpose/onnxdet.py

@@ -0,0 +1,128 @@
+# Code from the original DWPose Implementation: https://github.com/IDEA-Research/DWPose
+
+import cv2
+import numpy as np
+
+
+def nms(boxes, scores, nms_thr):
+    """Single class NMS implemented in Numpy."""
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= nms_thr)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+def multiclass_nms(boxes, scores, nms_thr, score_thr):
+    """Multiclass NMS implemented in Numpy. Class-aware version."""
+    final_dets = []
+    num_classes = scores.shape[1]
+    for cls_ind in range(num_classes):
+        cls_scores = scores[:, cls_ind]
+        valid_score_mask = cls_scores > score_thr
+        if valid_score_mask.sum() == 0:
+            continue
+        else:
+            valid_scores = cls_scores[valid_score_mask]
+            valid_boxes = boxes[valid_score_mask]
+            keep = nms(valid_boxes, valid_scores, nms_thr)
+            if len(keep) > 0:
+                cls_inds = np.ones((len(keep), 1)) * cls_ind
+                dets = np.concatenate([valid_boxes[keep], valid_scores[keep, None], cls_inds], 1)
+                final_dets.append(dets)
+    if len(final_dets) == 0:
+        return None
+    return np.concatenate(final_dets, 0)
+
+
+def demo_postprocess(outputs, img_size, p6=False):
+    grids = []
+    expanded_strides = []
+    strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
+
+    hsizes = [img_size[0] // stride for stride in strides]
+    wsizes = [img_size[1] // stride for stride in strides]
+
+    for hsize, wsize, stride in zip(hsizes, wsizes, strides, strict=False):
+        xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
+        grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
+        grids.append(grid)
+        shape = grid.shape[:2]
+        expanded_strides.append(np.full((*shape, 1), stride))
+
+    grids = np.concatenate(grids, 1)
+    expanded_strides = np.concatenate(expanded_strides, 1)
+    outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
+    outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
+
+    return outputs
+
+
+def preprocess(img, input_size, swap=(2, 0, 1)):
+    if len(img.shape) == 3:
+        padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
+    else:
+        padded_img = np.ones(input_size, dtype=np.uint8) * 114
+
+    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
+    resized_img = cv2.resize(
+        img,
+        (int(img.shape[1] * r), int(img.shape[0] * r)),
+        interpolation=cv2.INTER_LINEAR,
+    ).astype(np.uint8)
+    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
+
+    padded_img = padded_img.transpose(swap)
+    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
+    return padded_img, r
+
+
+def inference_detector(session, oriImg):
+    input_shape = (640, 640)
+    img, ratio = preprocess(oriImg, input_shape)
+
+    ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
+    output = session.run(None, ort_inputs)
+    predictions = demo_postprocess(output[0], input_shape)[0]
+
+    boxes = predictions[:, :4]
+    scores = predictions[:, 4:5] * predictions[:, 5:]
+
+    boxes_xyxy = np.ones_like(boxes)
+    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
+    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
+    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
+    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
+    boxes_xyxy /= ratio
+    dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
+    if dets is not None:
+        final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
+        isscore = final_scores > 0.3
+        iscat = final_cls_inds == 0
+        isbbox = [i and j for (i, j) in zip(isscore, iscat, strict=False)]
+        final_boxes = final_boxes[isbbox]
+    else:
+        final_boxes = np.array([])
+
+    return final_boxes

invokeai/backend/image_util/dw_openpose/onnxpose.py

@@ -0,0 +1,361 @@
+# Code from the original DWPose Implementation: https://github.com/IDEA-Research/DWPose
+
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+
+def preprocess(
+    img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Do preprocessing for RTMPose model inference.
+
+    Args:
+        img (np.ndarray): Input image in shape.
+        input_size (tuple): Input image size in shape (w, h).
+
+    Returns:
+        tuple:
+        - resized_img (np.ndarray): Preprocessed image.
+        - center (np.ndarray): Center of image.
+        - scale (np.ndarray): Scale of image.
+    """
+    # get shape of image
+    img_shape = img.shape[:2]
+    out_img, out_center, out_scale = [], [], []
+    if len(out_bbox) == 0:
+        out_bbox = [[0, 0, img_shape[1], img_shape[0]]]
+    for i in range(len(out_bbox)):
+        x0 = out_bbox[i][0]
+        y0 = out_bbox[i][1]
+        x1 = out_bbox[i][2]
+        y1 = out_bbox[i][3]
+        bbox = np.array([x0, y0, x1, y1])
+
+        # get center and scale
+        center, scale = bbox_xyxy2cs(bbox, padding=1.25)
+
+        # do affine transformation
+        resized_img, scale = top_down_affine(input_size, scale, center, img)
+
+        # normalize image
+        mean = np.array([123.675, 116.28, 103.53])
+        std = np.array([58.395, 57.12, 57.375])
+        resized_img = (resized_img - mean) / std
+
+        out_img.append(resized_img)
+        out_center.append(center)
+        out_scale.append(scale)
+
+    return out_img, out_center, out_scale
+
+
+def inference(sess: ort.InferenceSession, img: np.ndarray) -> np.ndarray:
+    """Inference RTMPose model.
+
+    Args:
+        sess (ort.InferenceSession): ONNXRuntime session.
+        img (np.ndarray): Input image in shape.
+
+    Returns:
+        outputs (np.ndarray): Output of RTMPose model.
+    """
+    all_out = []
+    # build input
+    for i in range(len(img)):
+        input = [img[i].transpose(2, 0, 1)]
+
+        # build output
+        sess_input = {sess.get_inputs()[0].name: input}
+        sess_output = []
+        for out in sess.get_outputs():
+            sess_output.append(out.name)
+
+        # run model
+        outputs = sess.run(sess_output, sess_input)
+        all_out.append(outputs)
+
+    return all_out
+
+
+def postprocess(
+    outputs: List[np.ndarray],
+    model_input_size: Tuple[int, int],
+    center: Tuple[int, int],
+    scale: Tuple[int, int],
+    simcc_split_ratio: float = 2.0,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Postprocess for RTMPose model output.
+
+    Args:
+        outputs (np.ndarray): Output of RTMPose model.
+        model_input_size (tuple): RTMPose model Input image size.
+        center (tuple): Center of bbox in shape (x, y).
+        scale (tuple): Scale of bbox in shape (w, h).
+        simcc_split_ratio (float): Split ratio of simcc.
+
+    Returns:
+        tuple:
+        - keypoints (np.ndarray): Rescaled keypoints.
+        - scores (np.ndarray): Model predict scores.
+    """
+    all_key = []
+    all_score = []
+    for i in range(len(outputs)):
+        # use simcc to decode
+        simcc_x, simcc_y = outputs[i]
+        keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio)
+
+        # rescale keypoints
+        keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2
+        all_key.append(keypoints[0])
+        all_score.append(scores[0])
+
+    return np.array(all_key), np.array(all_score)
+
+
+def bbox_xyxy2cs(bbox: np.ndarray, padding: float = 1.0) -> Tuple[np.ndarray, np.ndarray]:
+    """Transform the bbox format from (x,y,w,h) into (center, scale)
+
+    Args:
+        bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
+            as (left, top, right, bottom)
+        padding (float): BBox padding factor that will be multilied to scale.
+            Default: 1.0
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
+            (n, 2)
+        - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
+            (n, 2)
+    """
+    # convert single bbox from (4, ) to (1, 4)
+    dim = bbox.ndim
+    if dim == 1:
+        bbox = bbox[None, :]
+
+    # get bbox center and scale
+    x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
+    center = np.hstack([x1 + x2, y1 + y2]) * 0.5
+    scale = np.hstack([x2 - x1, y2 - y1]) * padding
+
+    if dim == 1:
+        center = center[0]
+        scale = scale[0]
+
+    return center, scale
+
+
+def _fix_aspect_ratio(bbox_scale: np.ndarray, aspect_ratio: float) -> np.ndarray:
+    """Extend the scale to match the given aspect ratio.
+
+    Args:
+        scale (np.ndarray): The image scale (w, h) in shape (2, )
+        aspect_ratio (float): The ratio of ``w/h``
+
+    Returns:
+        np.ndarray: The reshaped image scale in (2, )
+    """
+    w, h = np.hsplit(bbox_scale, [1])
+    bbox_scale = np.where(w > h * aspect_ratio, np.hstack([w, w / aspect_ratio]), np.hstack([h * aspect_ratio, h]))
+    return bbox_scale
+
+
+def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
+    """Rotate a point by an angle.
+
+    Args:
+        pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
+        angle_rad (float): rotation angle in radian
+
+    Returns:
+        np.ndarray: Rotated point in shape (2, )
+    """
+    sn, cs = np.sin(angle_rad), np.cos(angle_rad)
+    rot_mat = np.array([[cs, -sn], [sn, cs]])
+    return rot_mat @ pt
+
+
+def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    """To calculate the affine matrix, three pairs of points are required. This
+    function is used to get the 3rd point, given 2D points a & b.
+
+    The 3rd point is defined by rotating vector `a - b` by 90 degrees
+    anticlockwise, using b as the rotation center.
+
+    Args:
+        a (np.ndarray): The 1st point (x,y) in shape (2, )
+        b (np.ndarray): The 2nd point (x,y) in shape (2, )
+
+    Returns:
+        np.ndarray: The 3rd point.
+    """
+    direction = a - b
+    c = b + np.r_[-direction[1], direction[0]]
+    return c
+
+
+def get_warp_matrix(
+    center: np.ndarray,
+    scale: np.ndarray,
+    rot: float,
+    output_size: Tuple[int, int],
+    shift: Tuple[float, float] = (0.0, 0.0),
+    inv: bool = False,
+) -> np.ndarray:
+    """Calculate the affine transformation matrix that can warp the bbox area
+    in the input image to the output size.
+
+    Args:
+        center (np.ndarray[2, ]): Center of the bounding box (x, y).
+        scale (np.ndarray[2, ]): Scale of the bounding box
+            wrt [width, height].
+        rot (float): Rotation angle (degree).
+        output_size (np.ndarray[2, ] | list(2,)): Size of the
+            destination heatmaps.
+        shift (0-100%): Shift translation ratio wrt the width/height.
+            Default (0., 0.).
+        inv (bool): Option to inverse the affine transform direction.
+            (inv=False: src->dst or inv=True: dst->src)
+
+    Returns:
+        np.ndarray: A 2x3 transformation matrix
+    """
+    shift = np.array(shift)
+    src_w = scale[0]
+    dst_w = output_size[0]
+    dst_h = output_size[1]
+
+    # compute transformation matrix
+    rot_rad = np.deg2rad(rot)
+    src_dir = _rotate_point(np.array([0.0, src_w * -0.5]), rot_rad)
+    dst_dir = np.array([0.0, dst_w * -0.5])
+
+    # get four corners of the src rectangle in the original image
+    src = np.zeros((3, 2), dtype=np.float32)
+    src[0, :] = center + scale * shift
+    src[1, :] = center + src_dir + scale * shift
+    src[2, :] = _get_3rd_point(src[0, :], src[1, :])
+
+    # get four corners of the dst rectangle in the input image
+    dst = np.zeros((3, 2), dtype=np.float32)
+    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
+    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
+    dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
+
+    if inv:
+        warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+    else:
+        warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+    return warp_mat
+
+
+def top_down_affine(
+    input_size: dict, bbox_scale: dict, bbox_center: dict, img: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Get the bbox image as the model input by affine transform.
+
+    Args:
+        input_size (dict): The input size of the model.
+        bbox_scale (dict): The bbox scale of the img.
+        bbox_center (dict): The bbox center of the img.
+        img (np.ndarray): The original image.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: img after affine transform.
+        - np.ndarray[float32]: bbox scale after affine transform.
+    """
+    w, h = input_size
+    warp_size = (int(w), int(h))
+
+    # reshape bbox to fixed aspect ratio
+    bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h)
+
+    # get the affine matrix
+    center = bbox_center
+    scale = bbox_scale
+    rot = 0
+    warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
+
+    # do affine transform
+    img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
+
+    return img, bbox_scale
+
+
+def get_simcc_maximum(simcc_x: np.ndarray, simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """Get maximum response location and value from simcc representations.
+
+    Note:
+        instance number: N
+        num_keypoints: K
+        heatmap height: H
+        heatmap width: W
+
+    Args:
+        simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
+        simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
+
+    Returns:
+        tuple:
+        - locs (np.ndarray): locations of maximum heatmap responses in shape
+            (K, 2) or (N, K, 2)
+        - vals (np.ndarray): values of maximum heatmap responses in shape
+            (K,) or (N, K)
+    """
+    N, K, Wx = simcc_x.shape
+    simcc_x = simcc_x.reshape(N * K, -1)
+    simcc_y = simcc_y.reshape(N * K, -1)
+
+    # get maximum value locations
+    x_locs = np.argmax(simcc_x, axis=1)
+    y_locs = np.argmax(simcc_y, axis=1)
+    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
+    max_val_x = np.amax(simcc_x, axis=1)
+    max_val_y = np.amax(simcc_y, axis=1)
+
+    # get maximum value across x and y axis
+    mask = max_val_x > max_val_y
+    max_val_x[mask] = max_val_y[mask]
+    vals = max_val_x
+    locs[vals <= 0.0] = -1
+
+    # reshape
+    locs = locs.reshape(N, K, 2)
+    vals = vals.reshape(N, K)
+
+    return locs, vals
+
+
+def decode(simcc_x: np.ndarray, simcc_y: np.ndarray, simcc_split_ratio) -> Tuple[np.ndarray, np.ndarray]:
+    """Modulate simcc distribution with Gaussian.
+
+    Args:
+        simcc_x (np.ndarray[K, Wx]): model predicted simcc in x.
+        simcc_y (np.ndarray[K, Wy]): model predicted simcc in y.
+        simcc_split_ratio (int): The split ratio of simcc.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2)
+        - np.ndarray[float32]: scores in shape (K,) or (n, K)
+    """
+    keypoints, scores = get_simcc_maximum(simcc_x, simcc_y)
+    keypoints /= simcc_split_ratio
+
+    return keypoints, scores
+
+
+def inference_pose(session, out_bbox, oriImg):
+    h, w = session.get_inputs()[0].shape[2:]
+    model_input_size = (w, h)
+    resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size)
+    outputs = inference(session, resized_img)
+    keypoints, scores = postprocess(outputs, model_input_size, center, scale)
+
+    return keypoints, scores

invokeai/backend/image_util/dw_openpose/utils.py

@@ -0,0 +1,158 @@
+# Code from the original DWPose Implementation: https://github.com/IDEA-Research/DWPose
+
+import math
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+
+eps = 0.01
+NDArrayInt = npt.NDArray[np.uint8]
+
+
+def draw_bodypose(canvas: NDArrayInt, candidate: NDArrayInt, subset: NDArrayInt) -> NDArrayInt:
+    H, W, C = canvas.shape
+    candidate = np.array(candidate)
+    subset = np.array(subset)
+
+    stickwidth = 4
+
+    limbSeq = [
+        [2, 3],
+        [2, 6],
+        [3, 4],
+        [4, 5],
+        [6, 7],
+        [7, 8],
+        [2, 9],
+        [9, 10],
+        [10, 11],
+        [2, 12],
+        [12, 13],
+        [13, 14],
+        [2, 1],
+        [1, 15],
+        [15, 17],
+        [1, 16],
+        [16, 18],
+        [3, 17],
+        [6, 18],
+    ]
+
+    colors = [
+        [255, 0, 0],
+        [255, 85, 0],
+        [255, 170, 0],
+        [255, 255, 0],
+        [170, 255, 0],
+        [85, 255, 0],
+        [0, 255, 0],
+        [0, 255, 85],
+        [0, 255, 170],
+        [0, 255, 255],
+        [0, 170, 255],
+        [0, 85, 255],
+        [0, 0, 255],
+        [85, 0, 255],
+        [170, 0, 255],
+        [255, 0, 255],
+        [255, 0, 170],
+        [255, 0, 85],
+    ]
+
+    for i in range(17):
+        for n in range(len(subset)):
+            index = subset[n][np.array(limbSeq[i]) - 1]
+            if -1 in index:
+                continue
+            Y = candidate[index.astype(int), 0] * float(W)
+            X = candidate[index.astype(int), 1] * float(H)
+            mX = np.mean(X)
+            mY = np.mean(Y)
+            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+            polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
+            cv2.fillConvexPoly(canvas, polygon, colors[i])
+
+    canvas = (canvas * 0.6).astype(np.uint8)
+
+    for i in range(18):
+        for n in range(len(subset)):
+            index = int(subset[n][i])
+            if index == -1:
+                continue
+            x, y = candidate[index][0:2]
+            x = int(x * W)
+            y = int(y * H)
+            cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
+
+    return canvas
+
+
+def draw_handpose(canvas: NDArrayInt, all_hand_peaks: NDArrayInt) -> NDArrayInt:
+    H, W, C = canvas.shape
+
+    edges = [
+        [0, 1],
+        [1, 2],
+        [2, 3],
+        [3, 4],
+        [0, 5],
+        [5, 6],
+        [6, 7],
+        [7, 8],
+        [0, 9],
+        [9, 10],
+        [10, 11],
+        [11, 12],
+        [0, 13],
+        [13, 14],
+        [14, 15],
+        [15, 16],
+        [0, 17],
+        [17, 18],
+        [18, 19],
+        [19, 20],
+    ]
+
+    for peaks in all_hand_peaks:
+        peaks = np.array(peaks)
+
+        for ie, e in enumerate(edges):
+            x1, y1 = peaks[e[0]]
+            x2, y2 = peaks[e[1]]
+            x1 = int(x1 * W)
+            y1 = int(y1 * H)
+            x2 = int(x2 * W)
+            y2 = int(y2 * H)
+            if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
+                hsv_color = np.array([[[ie / float(len(edges)) * 180, 255, 255]]], dtype=np.uint8)
+                rgb_color = cv2.cvtColor(hsv_color, cv2.COLOR_HSV2RGB)[0, 0]
+                cv2.line(
+                    canvas,
+                    (x1, y1),
+                    (x2, y2),
+                    rgb_color.tolist(),
+                    thickness=2,
+                )
+
+        for _, keyponit in enumerate(peaks):
+            x, y = keyponit
+            x = int(x * W)
+            y = int(y * H)
+            if x > eps and y > eps:
+                cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
+    return canvas
+
+
+def draw_facepose(canvas: NDArrayInt, all_lmks: NDArrayInt) -> NDArrayInt:
+    H, W, C = canvas.shape
+    for lmks in all_lmks:
+        lmks = np.array(lmks)
+        for lmk in lmks:
+            x, y = lmk
+            x = int(x * W)
+            y = int(y * H)
+            if x > eps and y > eps:
+                cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1)
+    return canvas

invokeai/backend/image_util/grounding_dino/detection_result.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class BoundingBox(BaseModel):
+    """Bounding box helper class."""
+
+    xmin: int
+    ymin: int
+    xmax: int
+    ymax: int
+
+
+class DetectionResult(BaseModel):
+    """Detection result from Grounding DINO."""
+
+    score: float
+    label: str
+    box: BoundingBox
+    model_config = ConfigDict(
+        # Allow arbitrary types for mask, since it will be a numpy array.
+        arbitrary_types_allowed=True
+    )

invokeai/backend/image_util/grounding_dino/grounding_dino_pipeline.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+from PIL import Image
+from transformers.pipelines import ZeroShotObjectDetectionPipeline
+
+from invokeai.backend.image_util.grounding_dino.detection_result import DetectionResult
+from invokeai.backend.raw_model import RawModel
+
+
+class GroundingDinoPipeline(RawModel):
+    """A wrapper class for a ZeroShotObjectDetectionPipeline that makes it compatible with the model manager's memory
+    management system.
+    """
+
+    def __init__(self, pipeline: ZeroShotObjectDetectionPipeline):
+        self._pipeline = pipeline
+
+    def detect(self, image: Image.Image, candidate_labels: list[str], threshold: float = 0.1) -> list[DetectionResult]:
+        results = self._pipeline(image=image, candidate_labels=candidate_labels, threshold=threshold)
+        assert results is not None
+        results = [DetectionResult.model_validate(result) for result in results]
+        return results
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        # HACK(ryand): The GroundingDinoPipeline does not work on MPS devices. We only allow it to be moved to CPU or
+        # CUDA.
+        if device is not None and device.type not in {"cpu", "cuda"}:
+            device = None
+        self._pipeline.model.to(device=device, dtype=dtype)
+        self._pipeline.device = self._pipeline.model.device
+
+    def calc_size(self) -> int:
+        # HACK(ryand): Fix the circular import issue.
+        from invokeai.backend.model_manager.load.model_util import calc_module_size
+
+        return calc_module_size(self._pipeline.model)

invokeai/backend/image_util/hed.py

@@ -0,0 +1,217 @@
+# Adapted from https://github.com/huggingface/controlnet_aux
+
+import pathlib
+
+import cv2
+import huggingface_hub
+import numpy as np
+import torch
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from invokeai.backend.image_util.util import (
+    nms,
+    normalize_image_channel_count,
+    np_to_pil,
+    pil_to_np,
+    resize_image_to_resolution,
+    safe_step,
+)
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+class DoubleConvBlock(torch.nn.Module):
+    def __init__(self, input_channel, output_channel, layer_number):
+        super().__init__()
+        self.convs = torch.nn.Sequential()
+        self.convs.append(
+            torch.nn.Conv2d(
+                in_channels=input_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1
+            )
+        )
+        for _i in range(1, layer_number):
+            self.convs.append(
+                torch.nn.Conv2d(
+                    in_channels=output_channel,
+                    out_channels=output_channel,
+                    kernel_size=(3, 3),
+                    stride=(1, 1),
+                    padding=1,
+                )
+            )
+        self.projection = torch.nn.Conv2d(
+            in_channels=output_channel, out_channels=1, kernel_size=(1, 1), stride=(1, 1), padding=0
+        )
+
+    def __call__(self, x, down_sampling=False):
+        h = x
+        if down_sampling:
+            h = torch.nn.functional.max_pool2d(h, kernel_size=(2, 2), stride=(2, 2))
+        for conv in self.convs:
+            h = conv(h)
+            h = torch.nn.functional.relu(h)
+        return h, self.projection(h)
+
+
+class ControlNetHED_Apache2(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.norm = torch.nn.Parameter(torch.zeros(size=(1, 3, 1, 1)))
+        self.block1 = DoubleConvBlock(input_channel=3, output_channel=64, layer_number=2)
+        self.block2 = DoubleConvBlock(input_channel=64, output_channel=128, layer_number=2)
+        self.block3 = DoubleConvBlock(input_channel=128, output_channel=256, layer_number=3)
+        self.block4 = DoubleConvBlock(input_channel=256, output_channel=512, layer_number=3)
+        self.block5 = DoubleConvBlock(input_channel=512, output_channel=512, layer_number=3)
+
+    def __call__(self, x):
+        h = x - self.norm
+        h, projection1 = self.block1(h)
+        h, projection2 = self.block2(h, down_sampling=True)
+        h, projection3 = self.block3(h, down_sampling=True)
+        h, projection4 = self.block4(h, down_sampling=True)
+        h, projection5 = self.block5(h, down_sampling=True)
+        return projection1, projection2, projection3, projection4, projection5
+
+
+class HEDProcessor:
+    """Holistically-Nested Edge Detection.
+
+    On instantiation, loads the HED model from the HuggingFace Hub.
+    """
+
+    def __init__(self):
+        model_path = hf_hub_download("lllyasviel/Annotators", "ControlNetHED.pth")
+        self.network = ControlNetHED_Apache2()
+        self.network.load_state_dict(torch.load(model_path, map_location="cpu"))
+        self.network.float().eval()
+
+    def to(self, device: torch.device):
+        self.network.to(device)
+        return self
+
+    def run(
+        self,
+        input_image: Image.Image,
+        detect_resolution: int = 512,
+        image_resolution: int = 512,
+        safe: bool = False,
+        scribble: bool = False,
+    ) -> Image.Image:
+        """Processes an image and returns the detected edges.
+
+        Args:
+            input_image: The input image.
+            detect_resolution: The resolution to fit the image to before edge detection.
+            image_resolution: The resolution to fit the edges to before returning.
+            safe: Whether to apply safe step to the detected edges.
+            scribble: Whether to apply non-maximum suppression and Gaussian blur to the detected edges.
+
+        Returns:
+            The detected edges.
+        """
+        device = get_effective_device(self.network)
+        np_image = pil_to_np(input_image)
+        np_image = normalize_image_channel_count(np_image)
+        np_image = resize_image_to_resolution(np_image, detect_resolution)
+
+        assert np_image.ndim == 3
+        height, width, _channels = np_image.shape
+        with torch.no_grad():
+            image_hed = torch.from_numpy(np_image.copy()).float().to(device)
+            image_hed = rearrange(image_hed, "h w c -> 1 c h w")
+            edges = self.network(image_hed)
+            edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges]
+            edges = [cv2.resize(e, (width, height), interpolation=cv2.INTER_LINEAR) for e in edges]
+            edges = np.stack(edges, axis=2)
+            edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
+            if safe:
+                edge = safe_step(edge)
+            edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
+
+        detected_map = edge
+        detected_map = normalize_image_channel_count(detected_map)
+
+        img = resize_image_to_resolution(np_image, image_resolution)
+        height, width, _channels = img.shape
+
+        detected_map = cv2.resize(detected_map, (width, height), interpolation=cv2.INTER_LINEAR)
+
+        if scribble:
+            detected_map = nms(detected_map, 127, 3.0)
+            detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
+            detected_map[detected_map > 4] = 255
+            detected_map[detected_map < 255] = 0
+
+        return np_to_pil(detected_map)
+
+
+class HEDEdgeDetector:
+    """Simple wrapper around the HED model for detecting edges in an image."""
+
+    hf_repo_id = "lllyasviel/Annotators"
+    hf_filename = "ControlNetHED.pth"
+
+    def __init__(self, model: ControlNetHED_Apache2):
+        self.model = model
+
+    @classmethod
+    def get_model_url(cls) -> str:
+        """Get the URL to download the model from the Hugging Face Hub."""
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path) -> ControlNetHED_Apache2:
+        """Load the model from a file."""
+        model = ControlNetHED_Apache2()
+        model.load_state_dict(torch.load(model_path, map_location="cpu"))
+        model.float().eval()
+        return model
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, image: Image.Image, safe: bool = False, scribble: bool = False) -> Image.Image:
+        """Processes an image and returns the detected edges.
+
+        Args:
+            image: The input image.
+            safe: Whether to apply safe step to the detected edges.
+            scribble: Whether to apply non-maximum suppression and Gaussian blur to the detected edges.
+
+        Returns:
+            The detected edges.
+        """
+
+        device = get_effective_device(self.model)
+
+        np_image = pil_to_np(image)
+
+        height, width, _channels = np_image.shape
+
+        with torch.no_grad():
+            image_hed = torch.from_numpy(np_image.copy()).float().to(device)
+            image_hed = rearrange(image_hed, "h w c -> 1 c h w")
+            edges = self.model(image_hed)
+            edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges]
+            edges = [cv2.resize(e, (width, height), interpolation=cv2.INTER_LINEAR) for e in edges]
+            edges = np.stack(edges, axis=2)
+            edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
+            if safe:
+                edge = safe_step(edge)
+            edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
+
+        detected_map = edge
+
+        detected_map = cv2.resize(detected_map, (width, height), interpolation=cv2.INTER_LINEAR)
+
+        if scribble:
+            detected_map = nms(detected_map, 127, 3.0)
+            detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
+            detected_map[detected_map > 4] = 255
+            detected_map[detected_map < 255] = 0
+
+        output = np_to_pil(detected_map)
+
+        return output

invokeai/backend/image_util/imwatermark/vendor.py

@@ -0,0 +1,310 @@
+# This file is vendored from https://github.com/ShieldMnt/invisible-watermark
+#
+# `invisible-watermark` is MIT licensed as of August 23, 2025, when the code was copied into this repo.
+#
+# Why we vendored it in:
+# `invisible-watermark` has a dependency on `opencv-python`, which conflicts with Invoke's dependency on
+# `opencv-contrib-python`. It's easier to copy the code over than complicate the installation process by
+# requiring an extra post-install step of removing `opencv-python` and installing `opencv-contrib-python`.
+
+import struct
+import uuid
+import base64
+import cv2
+import numpy as np
+import pywt
+
+
+class WatermarkEncoder(object):
+    def __init__(self, content=b""):
+        seq = np.array([n for n in content], dtype=np.uint8)
+        self._watermarks = list(np.unpackbits(seq))
+        self._wmLen = len(self._watermarks)
+        self._wmType = "bytes"
+
+    def set_by_ipv4(self, addr):
+        bits = []
+        ips = addr.split(".")
+        for ip in ips:
+            bits += list(np.unpackbits(np.array([ip % 255], dtype=np.uint8)))
+        self._watermarks = bits
+        self._wmLen = len(self._watermarks)
+        self._wmType = "ipv4"
+        assert self._wmLen == 32
+
+    def set_by_uuid(self, uid):
+        u = uuid.UUID(uid)
+        self._wmType = "uuid"
+        seq = np.array([n for n in u.bytes], dtype=np.uint8)
+        self._watermarks = list(np.unpackbits(seq))
+        self._wmLen = len(self._watermarks)
+
+    def set_by_bytes(self, content):
+        self._wmType = "bytes"
+        seq = np.array([n for n in content], dtype=np.uint8)
+        self._watermarks = list(np.unpackbits(seq))
+        self._wmLen = len(self._watermarks)
+
+    def set_by_b16(self, b16):
+        content = base64.b16decode(b16)
+        self.set_by_bytes(content)
+        self._wmType = "b16"
+
+    def set_by_bits(self, bits=None):
+        if bits is None:
+            bits = []
+        self._watermarks = [int(bit) % 2 for bit in bits]
+        self._wmLen = len(self._watermarks)
+        self._wmType = "bits"
+
+    def set_watermark(self, wmType="bytes", content=""):
+        if wmType == "ipv4":
+            self.set_by_ipv4(content)
+        elif wmType == "uuid":
+            self.set_by_uuid(content)
+        elif wmType == "bits":
+            self.set_by_bits(content)
+        elif wmType == "bytes":
+            self.set_by_bytes(content)
+        elif wmType == "b16":
+            self.set_by_b16(content)
+        else:
+            raise NameError("%s is not supported" % wmType)
+
+    def get_length(self):
+        return self._wmLen
+
+    # @classmethod
+    # def loadModel(cls):
+    #     RivaWatermark.loadModel()
+
+    def encode(self, cv2Image, method="dwtDct", **configs):
+        (r, c, channels) = cv2Image.shape
+        if r * c < 256 * 256:
+            raise RuntimeError("image too small, should be larger than 256x256")
+
+        if method == "dwtDct":
+            embed = EmbedMaxDct(self._watermarks, wmLen=self._wmLen, **configs)
+            return embed.encode(cv2Image)
+        # elif method == 'dwtDctSvd':
+        #     embed = EmbedDwtDctSvd(self._watermarks, wmLen=self._wmLen, **configs)
+        #     return embed.encode(cv2Image)
+        # elif method == 'rivaGan':
+        #     embed = RivaWatermark(self._watermarks, self._wmLen)
+        #     return embed.encode(cv2Image)
+        else:
+            raise NameError("%s is not supported" % method)
+
+
+class WatermarkDecoder(object):
+    def __init__(self, wm_type="bytes", length=0):
+        self._wmType = wm_type
+        if wm_type == "ipv4":
+            self._wmLen = 32
+        elif wm_type == "uuid":
+            self._wmLen = 128
+        elif wm_type == "bytes":
+            self._wmLen = length
+        elif wm_type == "bits":
+            self._wmLen = length
+        elif wm_type == "b16":
+            self._wmLen = length
+        else:
+            raise NameError("%s is unsupported" % wm_type)
+
+    def reconstruct_ipv4(self, bits):
+        ips = [str(ip) for ip in list(np.packbits(bits))]
+        return ".".join(ips)
+
+    def reconstruct_uuid(self, bits):
+        nums = np.packbits(bits)
+        bstr = b""
+        for i in range(16):
+            bstr += struct.pack(">B", nums[i])
+
+        return str(uuid.UUID(bytes=bstr))
+
+    def reconstruct_bits(self, bits):
+        # return ''.join([str(b) for b in bits])
+        return bits
+
+    def reconstruct_b16(self, bits):
+        bstr = self.reconstruct_bytes(bits)
+        return base64.b16encode(bstr)
+
+    def reconstruct_bytes(self, bits):
+        nums = np.packbits(bits)
+        bstr = b""
+        for i in range(self._wmLen // 8):
+            bstr += struct.pack(">B", nums[i])
+        return bstr
+
+    def reconstruct(self, bits):
+        if len(bits) != self._wmLen:
+            raise RuntimeError("bits are not matched with watermark length")
+
+        if self._wmType == "ipv4":
+            return self.reconstruct_ipv4(bits)
+        elif self._wmType == "uuid":
+            return self.reconstruct_uuid(bits)
+        elif self._wmType == "bits":
+            return self.reconstruct_bits(bits)
+        elif self._wmType == "b16":
+            return self.reconstruct_b16(bits)
+        else:
+            return self.reconstruct_bytes(bits)
+
+    def decode(self, cv2Image, method="dwtDct", **configs):
+        (r, c, channels) = cv2Image.shape
+        if r * c < 256 * 256:
+            raise RuntimeError("image too small, should be larger than 256x256")
+
+        bits = []
+        if method == "dwtDct":
+            embed = EmbedMaxDct(watermarks=[], wmLen=self._wmLen, **configs)
+            bits = embed.decode(cv2Image)
+        # elif method == 'dwtDctSvd':
+        #     embed = EmbedDwtDctSvd(watermarks=[], wmLen=self._wmLen, **configs)
+        #     bits = embed.decode(cv2Image)
+        # elif method == 'rivaGan':
+        #     embed = RivaWatermark(watermarks=[], wmLen=self._wmLen, **configs)
+        #     bits = embed.decode(cv2Image)
+        else:
+            raise NameError("%s is not supported" % method)
+        return self.reconstruct(bits)
+
+    # @classmethod
+    # def loadModel(cls):
+    #     RivaWatermark.loadModel()
+
+
+class EmbedMaxDct(object):
+    def __init__(self, watermarks=None, wmLen=8, scales=None, block=4):
+        if watermarks is None:
+            watermarks = []
+        if scales is None:
+            scales = [0, 36, 36]
+        self._watermarks = watermarks
+        self._wmLen = wmLen
+        self._scales = scales
+        self._block = block
+
+    def encode(self, bgr):
+        (row, col, channels) = bgr.shape
+
+        yuv = cv2.cvtColor(bgr, cv2.COLOR_BGR2YUV)
+
+        for channel in range(2):
+            if self._scales[channel] <= 0:
+                continue
+
+            ca1, (h1, v1, d1) = pywt.dwt2(yuv[: row // 4 * 4, : col // 4 * 4, channel], "haar")
+            self.encode_frame(ca1, self._scales[channel])
+
+            yuv[: row // 4 * 4, : col // 4 * 4, channel] = pywt.idwt2((ca1, (v1, h1, d1)), "haar")
+
+        bgr_encoded = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
+        return bgr_encoded
+
+    def decode(self, bgr):
+        (row, col, channels) = bgr.shape
+
+        yuv = cv2.cvtColor(bgr, cv2.COLOR_BGR2YUV)
+
+        scores = [[] for i in range(self._wmLen)]
+        for channel in range(2):
+            if self._scales[channel] <= 0:
+                continue
+
+            ca1, (h1, v1, d1) = pywt.dwt2(yuv[: row // 4 * 4, : col // 4 * 4, channel], "haar")
+
+            scores = self.decode_frame(ca1, self._scales[channel], scores)
+
+        avgScores = list(map(lambda l: np.array(l).mean(), scores))
+
+        bits = np.array(avgScores) * 255 > 127
+        return bits
+
+    def decode_frame(self, frame, scale, scores):
+        (row, col) = frame.shape
+        num = 0
+
+        for i in range(row // self._block):
+            for j in range(col // self._block):
+                block = frame[
+                    i * self._block : i * self._block + self._block, j * self._block : j * self._block + self._block
+                ]
+
+                score = self.infer_dct_matrix(block, scale)
+                # score = self.infer_dct_svd(block, scale)
+                wmBit = num % self._wmLen
+                scores[wmBit].append(score)
+                num = num + 1
+
+        return scores
+
+    def diffuse_dct_svd(self, block, wmBit, scale):
+        u, s, v = np.linalg.svd(cv2.dct(block))
+
+        s[0] = (s[0] // scale + 0.25 + 0.5 * wmBit) * scale
+        return cv2.idct(np.dot(u, np.dot(np.diag(s), v)))
+
+    def infer_dct_svd(self, block, scale):
+        u, s, v = np.linalg.svd(cv2.dct(block))
+
+        score = 0
+        score = int((s[0] % scale) > scale * 0.5)
+        return score
+        if score >= 0.5:
+            return 1.0
+        else:
+            return 0.0
+
+    def diffuse_dct_matrix(self, block, wmBit, scale):
+        pos = np.argmax(abs(block.flatten()[1:])) + 1
+        i, j = pos // self._block, pos % self._block
+        val = block[i][j]
+        if val >= 0.0:
+            block[i][j] = (val // scale + 0.25 + 0.5 * wmBit) * scale
+        else:
+            val = abs(val)
+            block[i][j] = -1.0 * (val // scale + 0.25 + 0.5 * wmBit) * scale
+        return block
+
+    def infer_dct_matrix(self, block, scale):
+        pos = np.argmax(abs(block.flatten()[1:])) + 1
+        i, j = pos // self._block, pos % self._block
+
+        val = block[i][j]
+        if val < 0:
+            val = abs(val)
+
+        if (val % scale) > 0.5 * scale:
+            return 1
+        else:
+            return 0
+
+    def encode_frame(self, frame, scale):
+        """
+        frame is a matrix (M, N)
+
+        we get K (watermark bits size) blocks (self._block x self._block)
+
+        For i-th block, we encode watermark[i] bit into it
+        """
+        (row, col) = frame.shape
+        num = 0
+        for i in range(row // self._block):
+            for j in range(col // self._block):
+                block = frame[
+                    i * self._block : i * self._block + self._block, j * self._block : j * self._block + self._block
+                ]
+                wmBit = self._watermarks[(num % self._wmLen)]
+
+                diffusedBlock = self.diffuse_dct_matrix(block, wmBit, scale)
+                # diffusedBlock = self.diffuse_dct_svd(block, wmBit, scale)
+                frame[
+                    i * self._block : i * self._block + self._block, j * self._block : j * self._block + self._block
+                ] = diffusedBlock
+
+                num = num + 1

invokeai/backend/image_util/infill_methods/cv2_inpaint.py

@@ -0,0 +1,20 @@
+import cv2
+import numpy as np
+from PIL import Image
+
+
+def cv2_inpaint(image: Image.Image) -> Image.Image:
+    # Prepare Image
+    image_array = np.array(image.convert("RGB"))
+    image_cv = cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR)
+
+    # Prepare Mask From Alpha Channel
+    mask = image.split()[3].convert("RGB")
+    mask_array = np.array(mask)
+    mask_cv = cv2.cvtColor(mask_array, cv2.COLOR_BGR2GRAY)
+    mask_inv = cv2.bitwise_not(mask_cv)
+
+    # Inpaint Image
+    inpainted_result = cv2.inpaint(image_cv, mask_inv, 3, cv2.INPAINT_TELEA)
+    inpainted_image = Image.fromarray(cv2.cvtColor(inpainted_result, cv2.COLOR_BGR2RGB))
+    return inpainted_image

invokeai/backend/image_util/infill_methods/lama.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+import invokeai.backend.util.logging as logger
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+from invokeai.backend.model_manager.taxonomy import AnyModel
+
+
+def norm_img(np_img):
+    if len(np_img.shape) == 2:
+        np_img = np_img[:, :, np.newaxis]
+    np_img = np.transpose(np_img, (2, 0, 1))
+    np_img = np_img.astype("float32") / 255
+    return np_img
+
+
+class LaMA:
+    def __init__(self, model: AnyModel):
+        self._model = model
+
+    def __call__(self, input_image: Image.Image, *args: Any, **kwds: Any) -> Any:
+        image = np.asarray(input_image.convert("RGB"))
+        image = norm_img(image)
+
+        mask = input_image.split()[-1]
+        mask = np.asarray(mask)
+        mask = np.invert(mask)
+        mask = norm_img(mask)
+        mask = (mask > 0) * 1
+
+        device = get_effective_device(self._model)
+        image = torch.from_numpy(image).unsqueeze(0).to(device)
+        mask = torch.from_numpy(mask).unsqueeze(0).to(device)
+
+        with torch.inference_mode():
+            infilled_image = self._model(image, mask)
+
+        infilled_image = infilled_image[0].permute(1, 2, 0).detach().cpu().numpy()
+        infilled_image = np.clip(infilled_image * 255, 0, 255).astype("uint8")
+        infilled_image = Image.fromarray(infilled_image)
+
+        return infilled_image
+
+    @staticmethod
+    def load_jit_model(url_or_path: str | Path, device: torch.device | str = "cpu") -> torch.nn.Module:
+        model_path = url_or_path
+        logger.info(f"Loading model from: {model_path}")
+        model: torch.nn.Module = torch.jit.load(model_path, map_location="cpu").to(device)  # type: ignore
+        model.eval()
+        return model

invokeai/backend/image_util/infill_methods/mosaic.py

@@ -0,0 +1,60 @@
+from typing import Tuple
+
+import numpy as np
+from PIL import Image
+
+
+def infill_mosaic(
+    image: Image.Image,
+    tile_shape: Tuple[int, int] = (64, 64),
+    min_color: Tuple[int, int, int, int] = (0, 0, 0, 0),
+    max_color: Tuple[int, int, int, int] = (255, 255, 255, 0),
+) -> Image.Image:
+    """
+    image:PIL - A PIL Image
+    tile_shape: Tuple[int,int] - Tile width & Tile Height
+    min_color: Tuple[int,int,int] - RGB values for the lowest color to clip to (0-255)
+    max_color: Tuple[int,int,int] - RGB values for the highest color to clip to (0-255)
+    """
+
+    np_image = np.array(image)  # Convert image to np array
+    alpha = np_image[:, :, 3]  # Get the mask from the alpha channel of the image
+    non_transparent_pixels = np_image[alpha != 0, :3]  # List of non-transparent pixels
+
+    # Create color tiles to paste in the empty areas of the image
+    tile_width, tile_height = tile_shape
+
+    # Clip the range of colors in the image to a particular spectrum only
+    r_min, g_min, b_min, _ = min_color
+    r_max, g_max, b_max, _ = max_color
+    non_transparent_pixels[:, 0] = np.clip(non_transparent_pixels[:, 0], r_min, r_max)
+    non_transparent_pixels[:, 1] = np.clip(non_transparent_pixels[:, 1], g_min, g_max)
+    non_transparent_pixels[:, 2] = np.clip(non_transparent_pixels[:, 2], b_min, b_max)
+
+    tiles = []
+    for _ in range(256):
+        color = non_transparent_pixels[np.random.randint(len(non_transparent_pixels))]
+        tile = np.zeros((tile_height, tile_width, 3), dtype=np.uint8)
+        tile[:, :] = color
+        tiles.append(tile)
+
+    # Fill the transparent area with tiles
+    filled_image = np.zeros((image.height, image.width, 3), dtype=np.uint8)
+
+    for x in range(image.width):
+        for y in range(image.height):
+            tile = tiles[np.random.randint(len(tiles))]
+            try:
+                filled_image[
+                    y - (y % tile_height) : y - (y % tile_height) + tile_height,
+                    x - (x % tile_width) : x - (x % tile_width) + tile_width,
+                ] = tile
+            except ValueError:
+                # Need to handle edge cases - literally
+                pass
+
+    filled_image = Image.fromarray(filled_image)  # Convert the filled tiles image to PIL
+    image = Image.composite(
+        image, filled_image, image.split()[-1]
+    )  # Composite the original image on top of the filled tiles
+    return image

invokeai/backend/image_util/infill_methods/patchmatch.py

@@ -0,0 +1,67 @@
+"""
+This module defines a singleton object, "patchmatch" that
+wraps the actual patchmatch object. It respects the global
+"try_patchmatch" attribute, so that patchmatch loading can
+be suppressed or deferred
+"""
+
+import numpy as np
+from PIL import Image
+
+import invokeai.backend.util.logging as logger
+from invokeai.app.services.config.config_default import get_config
+
+
+class PatchMatch:
+    """
+    Thin class wrapper around the patchmatch function.
+    """
+
+    patch_match = None
+    tried_load: bool = False
+
+    def __init__(self):
+        super().__init__()
+
+    @classmethod
+    def _load_patch_match(cls):
+        if cls.tried_load:
+            return
+        if get_config().patchmatch:
+            from patchmatch import patch_match as pm
+
+            if pm.patchmatch_available:
+                logger.info("Patchmatch initialized")
+                cls.patch_match = pm
+            else:
+                logger.info("Patchmatch not loaded (nonfatal)")
+        else:
+            logger.info("Patchmatch loading disabled")
+        cls.tried_load = True
+
+    @classmethod
+    def patchmatch_available(cls) -> bool:
+        cls._load_patch_match()
+        if not cls.patch_match:
+            return False
+        return cls.patch_match.patchmatch_available
+
+    @classmethod
+    def inpaint(cls, image: Image.Image) -> Image.Image:
+        if cls.patch_match is None or not cls.patchmatch_available():
+            return image
+
+        np_image = np.array(image)
+        mask = 255 - np_image[:, :, 3]
+        infilled = cls.patch_match.inpaint(np_image[:, :, :3], mask, patch_size=3)
+        return Image.fromarray(infilled, mode="RGB")
+
+
+def infill_patchmatch(image: Image.Image) -> Image.Image:
+    IS_PATCHMATCH_AVAILABLE = PatchMatch.patchmatch_available()
+
+    if not IS_PATCHMATCH_AVAILABLE:
+        logger.warning("PatchMatch is not available on this system")
+        return image
+
+    return PatchMatch.inpaint(image)

invokeai/backend/image_util/infill_methods/tile.ipynb

@@ -0,0 +1,95 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"Smoke test for the tile infill\"\"\"\n",
+    "\n",
+    "from pathlib import Path\n",
+    "from typing import Optional\n",
+    "from PIL import Image\n",
+    "from invokeai.backend.image_util.infill_methods.tile import infill_tile\n",
+    "\n",
+    "images: list[tuple[str, Image.Image]] = []\n",
+    "\n",
+    "for i in sorted(Path(\"./test_images/\").glob(\"*.webp\")):\n",
+    "    images.append((i.name, Image.open(i)))\n",
+    "    images.append((i.name, Image.open(i).transpose(Image.FLIP_LEFT_RIGHT)))\n",
+    "    images.append((i.name, Image.open(i).transpose(Image.FLIP_TOP_BOTTOM)))\n",
+    "    images.append((i.name, Image.open(i).resize((512, 512))))\n",
+    "    images.append((i.name, Image.open(i).resize((1234, 461))))\n",
+    "\n",
+    "outputs: list[tuple[str, Image.Image, Image.Image, Optional[Image.Image]]] = []\n",
+    "\n",
+    "for name, image in images:\n",
+    "    try:\n",
+    "        output = infill_tile(image, seed=0, tile_size=32)\n",
+    "        outputs.append((name, image, output.infilled, output.tile_image))\n",
+    "    except ValueError as e:\n",
+    "        print(f\"Skipping image {name}: {e}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Display the images in jupyter notebook\n",
+    "import matplotlib.pyplot as plt\n",
+    "from PIL import ImageOps\n",
+    "\n",
+    "fig, axes = plt.subplots(len(outputs), 3, figsize=(10, 3 * len(outputs)))\n",
+    "plt.subplots_adjust(hspace=0)\n",
+    "\n",
+    "for i, (name, original, infilled, tile_image) in enumerate(outputs):\n",
+    "    # Add a border to each image, helps to see the edges\n",
+    "    size = original.size\n",
+    "    original = ImageOps.expand(original, border=5, fill=\"red\")\n",
+    "    filled = ImageOps.expand(infilled, border=5, fill=\"red\")\n",
+    "    if tile_image:\n",
+    "        tile_image = ImageOps.expand(tile_image, border=5, fill=\"red\")\n",
+    "\n",
+    "    axes[i, 0].imshow(original)\n",
+    "    axes[i, 0].axis(\"off\")\n",
+    "    axes[i, 0].set_title(f\"Original ({name} - {size})\")\n",
+    "\n",
+    "    if tile_image:\n",
+    "        axes[i, 1].imshow(tile_image)\n",
+    "        axes[i, 1].axis(\"off\")\n",
+    "        axes[i, 1].set_title(\"Tile Image\")\n",
+    "    else:\n",
+    "        axes[i, 1].axis(\"off\")\n",
+    "        axes[i, 1].set_title(\"NO TILES GENERATED (NO TRANSPARENCY)\")\n",
+    "\n",
+    "    axes[i, 2].imshow(filled)\n",
+    "    axes[i, 2].axis(\"off\")\n",
+    "    axes[i, 2].set_title(\"Filled\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".invokeai",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

invokeai/backend/image_util/infill_methods/tile.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+from PIL import Image
+
+
+def create_tile_pool(img_array: np.ndarray, tile_size: tuple[int, int]) -> list[np.ndarray]:
+    """
+    Create a pool of tiles from non-transparent areas of the image by systematically walking through the image.
+
+    Args:
+        img_array: numpy array of the image.
+        tile_size: tuple (tile_width, tile_height) specifying the size of each tile.
+
+    Returns:
+        A list of numpy arrays, each representing a tile.
+    """
+    tiles: list[np.ndarray] = []
+    rows, cols = img_array.shape[:2]
+    tile_width, tile_height = tile_size
+
+    for y in range(0, rows - tile_height + 1, tile_height):
+        for x in range(0, cols - tile_width + 1, tile_width):
+            tile = img_array[y : y + tile_height, x : x + tile_width]
+            # Check if the image has an alpha channel and the tile is completely opaque
+            if img_array.shape[2] == 4 and np.all(tile[:, :, 3] == 255):
+                tiles.append(tile)
+            elif img_array.shape[2] == 3:  # If no alpha channel, append the tile
+                tiles.append(tile)
+
+    if not tiles:
+        raise ValueError(
+            "Not enough opaque pixels to generate any tiles. Use a smaller tile size or a different image."
+        )
+
+    return tiles
+
+
+def create_filled_image(
+    img_array: np.ndarray, tile_pool: list[np.ndarray], tile_size: tuple[int, int], seed: int
+) -> np.ndarray:
+    """
+    Create an image of the same dimensions as the original, filled entirely with tiles from the pool.
+
+    Args:
+        img_array: numpy array of the original image.
+        tile_pool: A list of numpy arrays, each representing a tile.
+        tile_size: tuple (tile_width, tile_height) specifying the size of each tile.
+
+    Returns:
+        A numpy array representing the filled image.
+    """
+
+    rows, cols, _ = img_array.shape
+    tile_width, tile_height = tile_size
+
+    # Prep an empty RGB image
+    filled_img_array = np.zeros((rows, cols, 3), dtype=img_array.dtype)
+
+    # Make the random tile selection reproducible
+    rng = np.random.default_rng(seed)
+
+    for y in range(0, rows, tile_height):
+        for x in range(0, cols, tile_width):
+            # Pick a random tile from the pool
+            tile = tile_pool[rng.integers(len(tile_pool))]
+
+            # Calculate the space available (may be less than tile size near the edges)
+            space_y = min(tile_height, rows - y)
+            space_x = min(tile_width, cols - x)
+
+            # Crop the tile if necessary to fit into the available space
+            cropped_tile = tile[:space_y, :space_x, :3]
+
+            # Fill the available space with the (possibly cropped) tile
+            filled_img_array[y : y + space_y, x : x + space_x, :3] = cropped_tile
+
+    return filled_img_array
+
+
+@dataclass
+class InfillTileOutput:
+    infilled: Image.Image
+    tile_image: Optional[Image.Image] = None
+
+
+def infill_tile(image_to_infill: Image.Image, seed: int, tile_size: int) -> InfillTileOutput:
+    """Infills an image with random tiles from the image itself.
+
+    If the image is not an RGBA image, it is returned untouched.
+
+    Args:
+        image: The image to infill.
+        tile_size: The size of the tiles to use for infilling.
+
+    Raises:
+        ValueError: If there are not enough opaque pixels to generate any tiles.
+    """
+
+    if image_to_infill.mode != "RGBA":
+        return InfillTileOutput(infilled=image_to_infill)
+
+    # Internally, we want a tuple of (tile_width, tile_height). In the future, the tile size can be any rectangle.
+    _tile_size = (tile_size, tile_size)
+    np_image = np.array(image_to_infill, dtype=np.uint8)
+
+    # Create the pool of tiles that we will use to infill
+    tile_pool = create_tile_pool(np_image, _tile_size)
+
+    # Create an image from the tiles, same size as the original
+    tile_np_image = create_filled_image(np_image, tile_pool, _tile_size, seed)
+
+    # Paste the OG image over the tile image, effectively infilling the area
+    tile_image = Image.fromarray(tile_np_image, "RGB")
+    infilled = tile_image.copy()
+    infilled.paste(image_to_infill, (0, 0), image_to_infill.split()[-1])
+
+    # I think we want this to be "RGBA"?
+    infilled.convert("RGBA")
+
+    return InfillTileOutput(infilled=infilled, tile_image=tile_image)

invokeai/backend/image_util/invisible_watermark.py

@@ -0,0 +1,49 @@
+"""
+This module defines a singleton object, "invisible_watermark" that
+wraps the invisible watermark model. It respects the global "invisible_watermark"
+configuration variable, that allows the watermarking to be supressed.
+"""
+
+import cv2
+import numpy as np
+from PIL import Image
+
+import invokeai.backend.util.logging as logger
+from invokeai.backend.image_util.imwatermark.vendor import WatermarkDecoder, WatermarkEncoder
+
+
+class InvisibleWatermark:
+    """
+    Wrapper around InvisibleWatermark module.
+    """
+
+    @classmethod
+    def add_watermark(cls, image: Image.Image, watermark_text: str) -> Image.Image:
+        logger.debug(f'Applying invisible watermark "{watermark_text}"')
+        bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
+        encoder = WatermarkEncoder()
+        encoder.set_watermark("bytes", watermark_text.encode("utf-8"))
+        bgr_encoded = encoder.encode(bgr, "dwtDct")
+        return Image.fromarray(cv2.cvtColor(bgr_encoded, cv2.COLOR_BGR2RGB)).convert("RGBA")
+
+    @classmethod
+    def decode_watermark(cls, image: Image.Image, length: int = 8) -> str:
+        """Attempt to decode an invisible watermark from an image.
+
+        Args:
+            image: The PIL Image to decode the watermark from.
+            length: The expected watermark length in bytes. Must match the length used when encoding.
+                The WatermarkDecoder requires the length in bits; this value is multiplied by 8 internally.
+
+        Returns:
+            The decoded watermark text, or an empty string if no watermark is detected or decoding fails.
+        """
+        logger.debug("Attempting to decode invisible watermark")
+        try:
+            bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
+            decoder = WatermarkDecoder("bytes", length * 8)
+            watermark_bytes = decoder.decode(bgr, "dwtDct")
+            return watermark_bytes.decode("utf-8", errors="ignore").rstrip("\x00")
+        except Exception:
+            logger.debug("Failed to decode invisible watermark")
+            return ""

invokeai/backend/image_util/lineart.py

@@ -0,0 +1,228 @@
+"""Adapted from https://github.com/huggingface/controlnet_aux (Apache-2.0 license)."""
+
+import pathlib
+
+import cv2
+import huggingface_hub
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from invokeai.backend.image_util.util import (
+    normalize_image_channel_count,
+    np_to_pil,
+    pil_to_np,
+    resize_image_to_resolution,
+)
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features):
+        super(ResidualBlock, self).__init__()
+
+        conv_block = [
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            nn.InstanceNorm2d(in_features),
+            nn.ReLU(inplace=True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            nn.InstanceNorm2d(in_features),
+        ]
+
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        return x + self.conv_block(x)
+
+
+class Generator(nn.Module):
+    def __init__(self, input_nc, output_nc, n_residual_blocks=9, sigmoid=True):
+        super(Generator, self).__init__()
+
+        # Initial convolution block
+        model0 = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 64, 7), nn.InstanceNorm2d(64), nn.ReLU(inplace=True)]
+        self.model0 = nn.Sequential(*model0)
+
+        # Downsampling
+        model1 = []
+        in_features = 64
+        out_features = in_features * 2
+        for _ in range(2):
+            model1 += [
+                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+                nn.InstanceNorm2d(out_features),
+                nn.ReLU(inplace=True),
+            ]
+            in_features = out_features
+            out_features = in_features * 2
+        self.model1 = nn.Sequential(*model1)
+
+        model2 = []
+        # Residual blocks
+        for _ in range(n_residual_blocks):
+            model2 += [ResidualBlock(in_features)]
+        self.model2 = nn.Sequential(*model2)
+
+        # Upsampling
+        model3 = []
+        out_features = in_features // 2
+        for _ in range(2):
+            model3 += [
+                nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
+                nn.InstanceNorm2d(out_features),
+                nn.ReLU(inplace=True),
+            ]
+            in_features = out_features
+            out_features = in_features // 2
+        self.model3 = nn.Sequential(*model3)
+
+        # Output layer
+        model4 = [nn.ReflectionPad2d(3), nn.Conv2d(64, output_nc, 7)]
+        if sigmoid:
+            model4 += [nn.Sigmoid()]
+
+        self.model4 = nn.Sequential(*model4)
+
+    def forward(self, x, cond=None):
+        out = self.model0(x)
+        out = self.model1(out)
+        out = self.model2(out)
+        out = self.model3(out)
+        out = self.model4(out)
+
+        return out
+
+
+class LineartProcessor:
+    """Processor for lineart detection."""
+
+    def __init__(self):
+        model_path = hf_hub_download("lllyasviel/Annotators", "sk_model.pth")
+        self.model = Generator(3, 1, 3)
+        self.model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+        self.model.eval()
+
+        coarse_model_path = hf_hub_download("lllyasviel/Annotators", "sk_model2.pth")
+        self.model_coarse = Generator(3, 1, 3)
+        self.model_coarse.load_state_dict(torch.load(coarse_model_path, map_location=torch.device("cpu")))
+        self.model_coarse.eval()
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        self.model_coarse.to(device)
+        return self
+
+    def run(
+        self, input_image: Image.Image, coarse: bool = False, detect_resolution: int = 512, image_resolution: int = 512
+    ) -> Image.Image:
+        """Processes an image to detect lineart.
+
+        Args:
+            input_image: The input image.
+            coarse: Whether to use the coarse model.
+            detect_resolution: The resolution to fit the image to before edge detection.
+            image_resolution: The resolution of the output image.
+
+        Returns:
+            The detected lineart.
+        """
+        device = get_effective_device(self.model)
+
+        np_image = pil_to_np(input_image)
+        np_image = normalize_image_channel_count(np_image)
+        np_image = resize_image_to_resolution(np_image, detect_resolution)
+
+        model = self.model_coarse if coarse else self.model
+        assert np_image.ndim == 3
+        image = np_image
+        with torch.no_grad():
+            image = torch.from_numpy(image).float().to(device)
+            image = image / 255.0
+            image = rearrange(image, "h w c -> 1 c h w")
+            line = model(image)[0][0]
+
+            line = line.cpu().numpy()
+            line = (line * 255.0).clip(0, 255).astype(np.uint8)
+
+        detected_map = line
+
+        detected_map = normalize_image_channel_count(detected_map)
+
+        img = resize_image_to_resolution(np_image, image_resolution)
+        H, W, C = img.shape
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+        detected_map = 255 - detected_map
+
+        return np_to_pil(detected_map)
+
+
+class LineartEdgeDetector:
+    """Simple wrapper around the fine and coarse lineart models for detecting edges in an image."""
+
+    hf_repo_id = "lllyasviel/Annotators"
+    hf_filename_fine = "sk_model.pth"
+    hf_filename_coarse = "sk_model2.pth"
+
+    @classmethod
+    def get_model_url(cls, coarse: bool = False) -> str:
+        """Get the URL to download the model from the Hugging Face Hub."""
+        if coarse:
+            return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename_coarse)
+        else:
+            return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename_fine)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path) -> Generator:
+        """Load the model from a file."""
+        model = Generator(3, 1, 3)
+        model.load_state_dict(torch.load(model_path, map_location="cpu"))
+        model.float().eval()
+        return model
+
+    def __init__(self, model: Generator) -> None:
+        self.model = model
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, image: Image.Image) -> Image.Image:
+        """Detects edges in the input image with the selected lineart model.
+
+        Args:
+            input: The input image.
+            coarse: Whether to use the coarse model.
+
+        Returns:
+            The detected edges.
+        """
+        device = get_effective_device(self.model)
+
+        np_image = pil_to_np(image)
+
+        with torch.no_grad():
+            np_image = torch.from_numpy(np_image).float().to(device)
+            np_image = np_image / 255.0
+            np_image = rearrange(np_image, "h w c -> 1 c h w")
+            line = self.model(np_image)[0][0]
+
+            line = line.cpu().numpy()
+            line = (line * 255.0).clip(0, 255).astype(np.uint8)
+
+        detected_map = 255 - line
+
+        # The lineart model often outputs a lot of almost-black noise. SD1.5 ControlNets seem to be OK with this, but
+        # SDXL ControlNets are not - they need a cleaner map. 12 was experimentally determined to be a good threshold,
+        # eliminating all the noise while keeping the actual edges. Other approaches to thresholding may be better,
+        # for example stretching the contrast or removing noise.
+        detected_map[detected_map < 12] = 0
+
+        output = np_to_pil(detected_map)
+
+        return output

invokeai/backend/image_util/lineart_anime.py

@@ -0,0 +1,274 @@
+"""Adapted from https://github.com/huggingface/controlnet_aux (Apache-2.0 license)."""
+
+import functools
+import pathlib
+from typing import Optional
+
+import cv2
+import huggingface_hub
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from invokeai.backend.image_util.util import (
+    normalize_image_channel_count,
+    np_to_pil,
+    pil_to_np,
+    resize_image_to_resolution,
+)
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+class UnetGenerator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(
+        self,
+        input_nc: int,
+        output_nc: int,
+        num_downs: int,
+        ngf: int = 64,
+        norm_layer=nn.BatchNorm2d,
+        use_dropout: bool = False,
+    ):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(UnetGenerator, self).__init__()
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True
+        )  # add the innermost layer
+        for _ in range(num_downs - 5):  # add intermediate layers with ngf * 8 filters
+            unet_block = UnetSkipConnectionBlock(
+                ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout
+            )
+        # gradually reduce the number of filters from ngf * 8 to ngf
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer
+        )
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer
+        )
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        self.model = UnetSkipConnectionBlock(
+            output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer
+        )  # add the outermost layer
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class UnetSkipConnectionBlock(nn.Module):
+    """Defines the Unet submodule with skip connection.
+    X -------------------identity----------------------
+    |-- downsampling -- |submodule| -- upsampling --|
+    """
+
+    def __init__(
+        self,
+        outer_nc: int,
+        inner_nc: int,
+        input_nc: Optional[int] = None,
+        submodule=None,
+        outermost: bool = False,
+        innermost: bool = False,
+        norm_layer=nn.BatchNorm2d,
+        use_dropout: bool = False,
+    ):
+        """Construct a Unet submodule with skip connections.
+        Parameters:
+            outer_nc (int) -- the number of filters in the outer conv layer
+            inner_nc (int) -- the number of filters in the inner conv layer
+            input_nc (int) -- the number of channels in input images/features
+            submodule (UnetSkipConnectionBlock) -- previously defined submodules
+            outermost (bool)    -- if this module is the outermost module
+            innermost (bool)    -- if this module is the innermost module
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+        """
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if isinstance(norm_layer, functools.partial):
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:  # add skip connections
+            return torch.cat([x, self.model(x)], 1)
+
+
+class LineartAnimeProcessor:
+    """Processes an image to detect lineart."""
+
+    def __init__(self):
+        model_path = hf_hub_download("lllyasviel/Annotators", "netG.pth")
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+        self.model = UnetGenerator(3, 1, 8, 64, norm_layer=norm_layer, use_dropout=False)
+        ckpt = torch.load(model_path)
+        for key in list(ckpt.keys()):
+            if "module." in key:
+                ckpt[key.replace("module.", "")] = ckpt[key]
+                del ckpt[key]
+        self.model.load_state_dict(ckpt)
+        self.model.eval()
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, input_image: Image.Image, detect_resolution: int = 512, image_resolution: int = 512) -> Image.Image:
+        """Processes an image to detect lineart.
+
+        Args:
+            input_image: The input image.
+            detect_resolution: The resolution to use for detection.
+            image_resolution: The resolution to use for the output image.
+
+        Returns:
+            The detected lineart.
+        """
+        device = get_effective_device(self.model)
+        np_image = pil_to_np(input_image)
+
+        np_image = normalize_image_channel_count(np_image)
+        np_image = resize_image_to_resolution(np_image, detect_resolution)
+
+        H, W, C = np_image.shape
+        Hn = 256 * int(np.ceil(float(H) / 256.0))
+        Wn = 256 * int(np.ceil(float(W) / 256.0))
+        img = cv2.resize(np_image, (Wn, Hn), interpolation=cv2.INTER_CUBIC)
+        with torch.no_grad():
+            image_feed = torch.from_numpy(img).float().to(device)
+            image_feed = image_feed / 127.5 - 1.0
+            image_feed = rearrange(image_feed, "h w c -> 1 c h w")
+
+            line = self.model(image_feed)[0, 0] * 127.5 + 127.5
+            line = line.cpu().numpy()
+
+            line = cv2.resize(line, (W, H), interpolation=cv2.INTER_CUBIC)
+            line = line.clip(0, 255).astype(np.uint8)
+
+        detected_map = line
+
+        detected_map = normalize_image_channel_count(detected_map)
+
+        img = resize_image_to_resolution(np_image, image_resolution)
+        H, W, C = img.shape
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+        detected_map = 255 - detected_map
+
+        return np_to_pil(detected_map)
+
+
+class LineartAnimeEdgeDetector:
+    """Simple wrapper around the Lineart Anime model for detecting edges in an image."""
+
+    hf_repo_id = "lllyasviel/Annotators"
+    hf_filename = "netG.pth"
+
+    @classmethod
+    def get_model_url(cls) -> str:
+        """Get the URL to download the model from the Hugging Face Hub."""
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path) -> UnetGenerator:
+        """Load the model from a file."""
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+        model = UnetGenerator(3, 1, 8, 64, norm_layer=norm_layer, use_dropout=False)
+        ckpt = torch.load(model_path)
+        for key in list(ckpt.keys()):
+            if "module." in key:
+                ckpt[key.replace("module.", "")] = ckpt[key]
+                del ckpt[key]
+        model.load_state_dict(ckpt)
+        model.eval()
+        return model
+
+    def __init__(self, model: UnetGenerator) -> None:
+        self.model = model
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, image: Image.Image) -> Image.Image:
+        """Processes an image and returns the detected edges."""
+        device = get_effective_device(self.model)
+
+        np_image = pil_to_np(image)
+
+        height, width, _channels = np_image.shape
+        new_height = 256 * int(np.ceil(float(height) / 256.0))
+        new_width = 256 * int(np.ceil(float(width) / 256.0))
+
+        resized_img = cv2.resize(np_image, (new_width, new_height), interpolation=cv2.INTER_CUBIC)
+
+        with torch.no_grad():
+            image_feed = torch.from_numpy(resized_img).float().to(device)
+            image_feed = image_feed / 127.5 - 1.0
+            image_feed = rearrange(image_feed, "h w c -> 1 c h w")
+
+            line = self.model(image_feed)[0, 0] * 127.5 + 127.5
+            line = line.cpu().numpy()
+
+            line = cv2.resize(line, (width, height), interpolation=cv2.INTER_CUBIC)
+            line = line.clip(0, 255).astype(np.uint8)
+
+        detected_map = 255 - line
+
+        # The lineart model often outputs a lot of almost-black noise. SD1.5 ControlNets seem to be OK with this, but
+        # SDXL ControlNets are not - they need a cleaner map. 12 was experimentally determined to be a good threshold,
+        # eliminating all the noise while keeping the actual edges. Other approaches to thresholding may be better,
+        # for example stretching the contrast or removing noise.
+        detected_map[detected_map < 12] = 0
+
+        output = np_to_pil(detected_map)
+
+        return output

invokeai/backend/image_util/mediapipe_face/__init__.py

@@ -0,0 +1,15 @@
+# Adapted from https://github.com/huggingface/controlnet_aux
+
+from PIL import Image
+
+from invokeai.backend.image_util.mediapipe_face.mediapipe_face_common import generate_annotation
+from invokeai.backend.image_util.util import np_to_pil, pil_to_np
+
+
+def detect_faces(image: Image.Image, max_faces: int = 1, min_confidence: float = 0.5) -> Image.Image:
+    """Detects faces in an image using MediaPipe."""
+
+    np_img = pil_to_np(image)
+    detected_map = generate_annotation(np_img, max_faces, min_confidence)
+    detected_map_pil = np_to_pil(detected_map)
+    return detected_map_pil

invokeai/backend/image_util/mediapipe_face/mediapipe_face_common.py

@@ -0,0 +1,149 @@
+from typing import Mapping
+
+import mediapipe as mp
+import numpy
+
+mp_drawing = mp.solutions.drawing_utils
+mp_drawing_styles = mp.solutions.drawing_styles
+mp_face_detection = mp.solutions.face_detection  # Only for counting faces.
+mp_face_mesh = mp.solutions.face_mesh
+mp_face_connections = mp.solutions.face_mesh_connections.FACEMESH_TESSELATION
+mp_hand_connections = mp.solutions.hands_connections.HAND_CONNECTIONS
+mp_body_connections = mp.solutions.pose_connections.POSE_CONNECTIONS
+
+DrawingSpec = mp.solutions.drawing_styles.DrawingSpec
+PoseLandmark = mp.solutions.drawing_styles.PoseLandmark
+
+min_face_size_pixels: int = 64
+f_thick = 2
+f_rad = 1
+right_iris_draw = DrawingSpec(color=(10, 200, 250), thickness=f_thick, circle_radius=f_rad)
+right_eye_draw = DrawingSpec(color=(10, 200, 180), thickness=f_thick, circle_radius=f_rad)
+right_eyebrow_draw = DrawingSpec(color=(10, 220, 180), thickness=f_thick, circle_radius=f_rad)
+left_iris_draw = DrawingSpec(color=(250, 200, 10), thickness=f_thick, circle_radius=f_rad)
+left_eye_draw = DrawingSpec(color=(180, 200, 10), thickness=f_thick, circle_radius=f_rad)
+left_eyebrow_draw = DrawingSpec(color=(180, 220, 10), thickness=f_thick, circle_radius=f_rad)
+mouth_draw = DrawingSpec(color=(10, 180, 10), thickness=f_thick, circle_radius=f_rad)
+head_draw = DrawingSpec(color=(10, 200, 10), thickness=f_thick, circle_radius=f_rad)
+
+# mp_face_mesh.FACEMESH_CONTOURS has all the items we care about.
+face_connection_spec = {}
+for edge in mp_face_mesh.FACEMESH_FACE_OVAL:
+    face_connection_spec[edge] = head_draw
+for edge in mp_face_mesh.FACEMESH_LEFT_EYE:
+    face_connection_spec[edge] = left_eye_draw
+for edge in mp_face_mesh.FACEMESH_LEFT_EYEBROW:
+    face_connection_spec[edge] = left_eyebrow_draw
+# for edge in mp_face_mesh.FACEMESH_LEFT_IRIS:
+#    face_connection_spec[edge] = left_iris_draw
+for edge in mp_face_mesh.FACEMESH_RIGHT_EYE:
+    face_connection_spec[edge] = right_eye_draw
+for edge in mp_face_mesh.FACEMESH_RIGHT_EYEBROW:
+    face_connection_spec[edge] = right_eyebrow_draw
+# for edge in mp_face_mesh.FACEMESH_RIGHT_IRIS:
+#    face_connection_spec[edge] = right_iris_draw
+for edge in mp_face_mesh.FACEMESH_LIPS:
+    face_connection_spec[edge] = mouth_draw
+iris_landmark_spec = {468: right_iris_draw, 473: left_iris_draw}
+
+
+def draw_pupils(image, landmark_list, drawing_spec, halfwidth: int = 2):
+    """We have a custom function to draw the pupils because the mp.draw_landmarks method requires a parameter for all
+    landmarks.  Until our PR is merged into mediapipe, we need this separate method."""
+    if len(image.shape) != 3:
+        raise ValueError("Input image must be H,W,C.")
+    image_rows, image_cols, image_channels = image.shape
+    if image_channels != 3:  # BGR channels
+        raise ValueError("Input image must contain three channel bgr data.")
+    for idx, landmark in enumerate(landmark_list.landmark):
+        if (landmark.HasField("visibility") and landmark.visibility < 0.9) or (
+            landmark.HasField("presence") and landmark.presence < 0.5
+        ):
+            continue
+        if landmark.x >= 1.0 or landmark.x < 0 or landmark.y >= 1.0 or landmark.y < 0:
+            continue
+        image_x = int(image_cols * landmark.x)
+        image_y = int(image_rows * landmark.y)
+        draw_color = None
+        if isinstance(drawing_spec, Mapping):
+            if drawing_spec.get(idx) is None:
+                continue
+            else:
+                draw_color = drawing_spec[idx].color
+        elif isinstance(drawing_spec, DrawingSpec):
+            draw_color = drawing_spec.color
+        image[image_y - halfwidth : image_y + halfwidth, image_x - halfwidth : image_x + halfwidth, :] = draw_color
+
+
+def reverse_channels(image):
+    """Given a numpy array in RGB form, convert to BGR.  Will also convert from BGR to RGB."""
+    # im[:,:,::-1] is a neat hack to convert BGR to RGB by reversing the indexing order.
+    # im[:,:,::[2,1,0]] would also work but makes a copy of the data.
+    return image[:, :, ::-1]
+
+
+def generate_annotation(img_rgb, max_faces: int, min_confidence: float):
+    """
+    Find up to 'max_faces' inside the provided input image.
+    If min_face_size_pixels is provided and nonzero it will be used to filter faces that occupy less than this many
+    pixels in the image.
+    """
+    with mp_face_mesh.FaceMesh(
+        static_image_mode=True,
+        max_num_faces=max_faces,
+        refine_landmarks=True,
+        min_detection_confidence=min_confidence,
+    ) as facemesh:
+        img_height, img_width, img_channels = img_rgb.shape
+        assert img_channels == 3
+
+        results = facemesh.process(img_rgb).multi_face_landmarks
+
+        if results is None:
+            print("No faces detected in controlnet image for Mediapipe face annotator.")
+            return numpy.zeros_like(img_rgb)
+
+        # Filter faces that are too small
+        filtered_landmarks = []
+        for lm in results:
+            landmarks = lm.landmark
+            face_rect = [
+                landmarks[0].x,
+                landmarks[0].y,
+                landmarks[0].x,
+                landmarks[0].y,
+            ]  # Left, up, right, down.
+            for i in range(len(landmarks)):
+                face_rect[0] = min(face_rect[0], landmarks[i].x)
+                face_rect[1] = min(face_rect[1], landmarks[i].y)
+                face_rect[2] = max(face_rect[2], landmarks[i].x)
+                face_rect[3] = max(face_rect[3], landmarks[i].y)
+            if min_face_size_pixels > 0:
+                face_width = abs(face_rect[2] - face_rect[0])
+                face_height = abs(face_rect[3] - face_rect[1])
+                face_width_pixels = face_width * img_width
+                face_height_pixels = face_height * img_height
+                face_size = min(face_width_pixels, face_height_pixels)
+                if face_size >= min_face_size_pixels:
+                    filtered_landmarks.append(lm)
+            else:
+                filtered_landmarks.append(lm)
+
+        # Annotations are drawn in BGR for some reason, but we don't need to flip a zero-filled image at the start.
+        empty = numpy.zeros_like(img_rgb)
+
+        # Draw detected faces:
+        for face_landmarks in filtered_landmarks:
+            mp_drawing.draw_landmarks(
+                empty,
+                face_landmarks,
+                connections=face_connection_spec.keys(),
+                landmark_drawing_spec=None,
+                connection_drawing_spec=face_connection_spec,
+            )
+            draw_pupils(empty, face_landmarks, iris_landmark_spec, 2)
+
+        # Flip BGR back to RGB.
+        empty = reverse_channels(empty).copy()
+
+        return empty

invokeai/backend/image_util/mlsd/__init__.py

@@ -0,0 +1,66 @@
+# Adapted from https://github.com/huggingface/controlnet_aux
+
+import pathlib
+
+import cv2
+import huggingface_hub
+import numpy as np
+import torch
+from PIL import Image
+
+from invokeai.backend.image_util.mlsd.models.mbv2_mlsd_large import MobileV2_MLSD_Large
+from invokeai.backend.image_util.mlsd.utils import pred_lines
+from invokeai.backend.image_util.util import np_to_pil, pil_to_np, resize_to_multiple
+
+
+class MLSDDetector:
+    """Simple wrapper around a MLSD model for detecting edges as line segments in an image."""
+
+    hf_repo_id = "lllyasviel/ControlNet"
+    hf_filename = "annotator/ckpts/mlsd_large_512_fp32.pth"
+
+    @classmethod
+    def get_model_url(cls) -> str:
+        """Get the URL to download the model from the Hugging Face Hub."""
+
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path) -> MobileV2_MLSD_Large:
+        """Load the model from a file."""
+
+        model = MobileV2_MLSD_Large()
+        model.load_state_dict(torch.load(model_path), strict=True)
+        model.eval()
+        return model
+
+    def __init__(self, model: MobileV2_MLSD_Large) -> None:
+        self.model = model
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, image: Image.Image, score_threshold: float = 0.1, distance_threshold: float = 20.0) -> Image.Image:
+        """Processes an image and returns the detected edges."""
+
+        np_img = pil_to_np(image)
+
+        height, width, _channels = np_img.shape
+
+        # This model requires the input image to have a resolution that is a multiple of 64
+        np_img = resize_to_multiple(np_img, 64)
+        img_output = np.zeros_like(np_img)
+
+        with torch.no_grad():
+            lines = pred_lines(np_img, self.model, [np_img.shape[0], np_img.shape[1]], score_threshold, distance_threshold)
+            for line in lines:
+                x_start, y_start, x_end, y_end = [int(val) for val in line]
+                cv2.line(img_output, (x_start, y_start), (x_end, y_end), [255, 255, 255], 1)
+
+        detected_map = img_output[:, :, 0]
+
+        # Back to the original size
+        output_image = cv2.resize(detected_map, (width, height), interpolation=cv2.INTER_LINEAR)
+
+        return np_to_pil(output_image)

invokeai/backend/image_util/mlsd/models/mbv2_mlsd_large.py

@@ -0,0 +1,290 @@
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from torch.nn import functional as F
+
+
+class BlockTypeA(nn.Module):
+    def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale = True):
+        super(BlockTypeA, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c2, out_c2, kernel_size=1),
+            nn.BatchNorm2d(out_c2),
+            nn.ReLU(inplace=True)
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c1, out_c1, kernel_size=1),
+            nn.BatchNorm2d(out_c1),
+            nn.ReLU(inplace=True)
+        )
+        self.upscale = upscale
+
+    def forward(self, a, b):
+        b = self.conv1(b)
+        a = self.conv2(a)
+        if self.upscale:
+             b = F.interpolate(b, scale_factor=2.0, mode='bilinear', align_corners=True)
+        return torch.cat((a, b), dim=1)
+
+
+class BlockTypeB(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeB, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_c),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        x = self.conv1(x) + x
+        x = self.conv2(x)
+        return x
+
+class BlockTypeC(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeC, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=5, dilation=5),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        return x
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = 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 < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class ConvBNReLU(nn.Sequential):
+    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
+        self.channel_pad = out_planes - in_planes
+        self.stride = stride
+        #padding = (kernel_size - 1) // 2
+
+        # TFLite uses slightly different padding than PyTorch
+        if stride == 2:
+            padding = 0
+        else:
+            padding = (kernel_size - 1) // 2
+
+        super(ConvBNReLU, self).__init__(
+            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU6(inplace=True)
+        )
+        self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
+
+
+    def forward(self, x):
+        # TFLite uses  different padding
+        if self.stride == 2:
+            x = F.pad(x, (0, 1, 0, 1), "constant", 0)
+            #print(x.shape)
+
+        for module in self:
+            if not isinstance(module, nn.MaxPool2d):
+                x = module(x)
+        return x
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = int(round(inp * expand_ratio))
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        layers = []
+        if expand_ratio != 1:
+            # pw
+            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
+        layers.extend([
+            # dw
+            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
+            # pw-linear
+            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(oup),
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, pretrained=True):
+        """
+        MobileNet V2 main class
+        Args:
+            num_classes (int): Number of classes
+            width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
+            inverted_residual_setting: Network structure
+            round_nearest (int): Round the number of channels in each layer to be a multiple of this number
+            Set to 1 to turn off rounding
+            block: Module specifying inverted residual building block for mobilenet
+        """
+        super(MobileNetV2, self).__init__()
+
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        width_mult = 1.0
+        round_nearest = 8
+
+        inverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            [6, 96, 3, 1],
+            #[6, 160, 3, 2],
+            #[6, 320, 1, 1],
+        ]
+
+        # only check the first element, assuming user knows t,c,n,s are required
+        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
+            raise ValueError("inverted_residual_setting should be non-empty "
+                             "or a 4-element list, got {}".format(inverted_residual_setting))
+
+        # building first layer
+        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
+        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
+        features = [ConvBNReLU(4, input_channel, stride=2)]
+        # building inverted residual blocks
+        for t, c, n, s in inverted_residual_setting:
+            output_channel = _make_divisible(c * width_mult, round_nearest)
+            for i in range(n):
+                stride = s if i == 0 else 1
+                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
+                input_channel = output_channel
+
+        self.features = nn.Sequential(*features)
+        self.fpn_selected = [1, 3, 6, 10, 13]
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.zeros_(m.bias)
+        if pretrained:
+           self._load_pretrained_model()
+
+    def _forward_impl(self, x):
+        # This exists since TorchScript doesn't support inheritance, so the superclass method
+        # (this one) needs to have a name other than `forward` that can be accessed in a subclass
+        fpn_features = []
+        for i, f in enumerate(self.features):
+            if i > self.fpn_selected[-1]:
+                break
+            x = f(x)
+            if i in self.fpn_selected:
+                fpn_features.append(x)
+
+        c1, c2, c3, c4, c5 = fpn_features
+        return c1, c2, c3, c4, c5
+
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+    def _load_pretrained_model(self):
+        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+
+class MobileV2_MLSD_Large(nn.Module):
+    def __init__(self):
+        super(MobileV2_MLSD_Large, self).__init__()
+
+        self.backbone = MobileNetV2(pretrained=False)
+        ## A, B
+        self.block15 = BlockTypeA(in_c1= 64, in_c2= 96,
+                                  out_c1= 64, out_c2=64,
+                                  upscale=False)
+        self.block16 = BlockTypeB(128, 64)
+
+        ## A, B
+        self.block17 = BlockTypeA(in_c1 = 32,  in_c2 = 64,
+                                  out_c1= 64,  out_c2= 64)
+        self.block18 = BlockTypeB(128, 64)
+
+        ## A, B
+        self.block19 = BlockTypeA(in_c1=24, in_c2=64,
+                                  out_c1=64, out_c2=64)
+        self.block20 = BlockTypeB(128, 64)
+
+        ## A, B, C
+        self.block21 = BlockTypeA(in_c1=16, in_c2=64,
+                                  out_c1=64, out_c2=64)
+        self.block22 = BlockTypeB(128, 64)
+
+        self.block23 = BlockTypeC(64, 16)
+
+    def forward(self, x):
+        c1, c2, c3, c4, c5 = self.backbone(x)
+
+        x = self.block15(c4, c5)
+        x = self.block16(x)
+
+        x = self.block17(c3, x)
+        x = self.block18(x)
+
+        x = self.block19(c2, x)
+        x = self.block20(x)
+
+        x = self.block21(c1, x)
+        x = self.block22(x)
+        x = self.block23(x)
+        x = x[:, 7:, :, :]
+
+        return x

invokeai/backend/image_util/mlsd/models/mbv2_mlsd_tiny.py

@@ -0,0 +1,273 @@
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from torch.nn import functional as F
+
+
+class BlockTypeA(nn.Module):
+    def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale = True):
+        super(BlockTypeA, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c2, out_c2, kernel_size=1),
+            nn.BatchNorm2d(out_c2),
+            nn.ReLU(inplace=True)
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c1, out_c1, kernel_size=1),
+            nn.BatchNorm2d(out_c1),
+            nn.ReLU(inplace=True)
+        )
+        self.upscale = upscale
+
+    def forward(self, a, b):
+        b = self.conv1(b)
+        a = self.conv2(a)
+        b = F.interpolate(b, scale_factor=2.0, mode='bilinear', align_corners=True)
+        return torch.cat((a, b), dim=1)
+
+
+class BlockTypeB(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeB, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_c),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        x = self.conv1(x) + x
+        x = self.conv2(x)
+        return x
+
+class BlockTypeC(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeC, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=5, dilation=5),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        return x
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = 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 < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class ConvBNReLU(nn.Sequential):
+    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
+        self.channel_pad = out_planes - in_planes
+        self.stride = stride
+        #padding = (kernel_size - 1) // 2
+
+        # TFLite uses slightly different padding than PyTorch
+        if stride == 2:
+            padding = 0
+        else:
+            padding = (kernel_size - 1) // 2
+
+        super(ConvBNReLU, self).__init__(
+            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU6(inplace=True)
+        )
+        self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
+
+
+    def forward(self, x):
+        # TFLite uses  different padding
+        if self.stride == 2:
+            x = F.pad(x, (0, 1, 0, 1), "constant", 0)
+            #print(x.shape)
+
+        for module in self:
+            if not isinstance(module, nn.MaxPool2d):
+                x = module(x)
+        return x
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = int(round(inp * expand_ratio))
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        layers = []
+        if expand_ratio != 1:
+            # pw
+            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
+        layers.extend([
+            # dw
+            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
+            # pw-linear
+            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(oup),
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, pretrained=True):
+        """
+        MobileNet V2 main class
+        Args:
+            num_classes (int): Number of classes
+            width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
+            inverted_residual_setting: Network structure
+            round_nearest (int): Round the number of channels in each layer to be a multiple of this number
+            Set to 1 to turn off rounding
+            block: Module specifying inverted residual building block for mobilenet
+        """
+        super(MobileNetV2, self).__init__()
+
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        width_mult = 1.0
+        round_nearest = 8
+
+        inverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            #[6, 96, 3, 1],
+            #[6, 160, 3, 2],
+            #[6, 320, 1, 1],
+        ]
+
+        # only check the first element, assuming user knows t,c,n,s are required
+        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
+            raise ValueError("inverted_residual_setting should be non-empty "
+                             "or a 4-element list, got {}".format(inverted_residual_setting))
+
+        # building first layer
+        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
+        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
+        features = [ConvBNReLU(4, input_channel, stride=2)]
+        # building inverted residual blocks
+        for t, c, n, s in inverted_residual_setting:
+            output_channel = _make_divisible(c * width_mult, round_nearest)
+            for i in range(n):
+                stride = s if i == 0 else 1
+                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
+                input_channel = output_channel
+        self.features = nn.Sequential(*features)
+
+        self.fpn_selected = [3, 6, 10]
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.zeros_(m.bias)
+
+        #if pretrained:
+        #    self._load_pretrained_model()
+
+    def _forward_impl(self, x):
+        # This exists since TorchScript doesn't support inheritance, so the superclass method
+        # (this one) needs to have a name other than `forward` that can be accessed in a subclass
+        fpn_features = []
+        for i, f in enumerate(self.features):
+            if i > self.fpn_selected[-1]:
+                break
+            x = f(x)
+            if i in self.fpn_selected:
+                fpn_features.append(x)
+
+        c2, c3, c4 = fpn_features
+        return c2, c3, c4
+
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+    def _load_pretrained_model(self):
+        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+
+class MobileV2_MLSD_Tiny(nn.Module):
+    def __init__(self):
+        super(MobileV2_MLSD_Tiny, self).__init__()
+
+        self.backbone = MobileNetV2(pretrained=True)
+
+        self.block12 = BlockTypeA(in_c1= 32, in_c2= 64,
+                                  out_c1= 64, out_c2=64)
+        self.block13 = BlockTypeB(128, 64)
+
+        self.block14 = BlockTypeA(in_c1 = 24,  in_c2 = 64,
+                                  out_c1= 32,  out_c2= 32)
+        self.block15 = BlockTypeB(64, 64)
+
+        self.block16 = BlockTypeC(64, 16)
+
+    def forward(self, x):
+        c2, c3, c4 = self.backbone(x)
+
+        x = self.block12(c3, c4)
+        x = self.block13(x)
+        x = self.block14(c2, x)
+        x = self.block15(x)
+        x = self.block16(x)
+        x = x[:, 7:, :, :]
+        #print(x.shape)
+        x = F.interpolate(x, scale_factor=2.0, mode='bilinear', align_corners=True)
+
+        return x

invokeai/backend/image_util/mlsd/utils.py

@@ -0,0 +1,589 @@
+'''
+modified by  lihaoweicv
+pytorch version
+'''
+
+'''
+M-LSD
+Copyright 2021-present NAVER Corp.
+Apache License v2.0
+'''
+
+import cv2
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+def deccode_output_score_and_ptss(tpMap, topk_n = 200, ksize = 5):
+    '''
+    tpMap:
+    center: tpMap[1, 0, :, :]
+    displacement: tpMap[1, 1:5, :, :]
+    '''
+    b, c, h, w = tpMap.shape
+    assert  b==1, 'only support bsize==1'
+    displacement = tpMap[:, 1:5, :, :][0]
+    center = tpMap[:, 0, :, :]
+    heat = torch.sigmoid(center)
+    hmax = F.max_pool2d( heat, (ksize, ksize), stride=1, padding=(ksize-1)//2)
+    keep = (hmax == heat).float()
+    heat = heat * keep
+    heat = heat.reshape(-1, )
+
+    scores, indices = torch.topk(heat, topk_n, dim=-1, largest=True)
+    yy = torch.floor_divide(indices, w).unsqueeze(-1)
+    xx = torch.fmod(indices, w).unsqueeze(-1)
+    ptss = torch.cat((yy, xx),dim=-1)
+
+    ptss   = ptss.detach().cpu().numpy()
+    scores = scores.detach().cpu().numpy()
+    displacement = displacement.detach().cpu().numpy()
+    displacement = displacement.transpose((1,2,0))
+    return  ptss, scores, displacement
+
+
+def pred_lines(image, model,
+               input_shape=[512, 512],
+               score_thr=0.10,
+               dist_thr=20.0):
+    h, w, _ = image.shape
+
+    device = get_effective_device(model)
+    h_ratio, w_ratio = [h / input_shape[0], w / input_shape[1]]
+
+    resized_image = np.concatenate([cv2.resize(image, (input_shape[1], input_shape[0]), interpolation=cv2.INTER_AREA),
+                                    np.ones([input_shape[0], input_shape[1], 1])], axis=-1)
+
+    resized_image = resized_image.transpose((2,0,1))
+    batch_image = np.expand_dims(resized_image, axis=0).astype('float32')
+    batch_image = (batch_image / 127.5) - 1.0
+
+    batch_image = torch.from_numpy(batch_image).float()
+    batch_image = batch_image.to(device)
+    outputs = model(batch_image)
+    pts, pts_score, vmap = deccode_output_score_and_ptss(outputs, 200, 3)
+    start = vmap[:, :, :2]
+    end = vmap[:, :, 2:]
+    dist_map = np.sqrt(np.sum((start - end) ** 2, axis=-1))
+
+    segments_list = []
+    for center, score in zip(pts, pts_score, strict=False):
+        y, x = center
+        distance = dist_map[y, x]
+        if score > score_thr and distance > dist_thr:
+            disp_x_start, disp_y_start, disp_x_end, disp_y_end = vmap[y, x, :]
+            x_start = x + disp_x_start
+            y_start = y + disp_y_start
+            x_end = x + disp_x_end
+            y_end = y + disp_y_end
+            segments_list.append([x_start, y_start, x_end, y_end])
+
+    if segments_list:
+        lines = 2 * np.array(segments_list)  # 256 > 512
+        lines[:, 0] = lines[:, 0] * w_ratio
+        lines[:, 1] = lines[:, 1] * h_ratio
+        lines[:, 2] = lines[:, 2] * w_ratio
+        lines[:, 3] = lines[:, 3] * h_ratio
+    else:
+        # No segments detected - return empty array
+        lines = np.array([])
+
+    return lines
+
+
+def pred_squares(image,
+                 model,
+                 input_shape=[512, 512],
+                 params={'score': 0.06,
+                         'outside_ratio': 0.28,
+                         'inside_ratio': 0.45,
+                         'w_overlap': 0.0,
+                         'w_degree': 1.95,
+                         'w_length': 0.0,
+                         'w_area': 1.86,
+                         'w_center': 0.14}):
+    '''
+    shape = [height, width]
+    '''
+    h, w, _ = image.shape
+    original_shape = [h, w]
+    device = get_effective_device(model)
+
+    resized_image = np.concatenate([cv2.resize(image, (input_shape[0], input_shape[1]), interpolation=cv2.INTER_AREA),
+                                    np.ones([input_shape[0], input_shape[1], 1])], axis=-1)
+    resized_image = resized_image.transpose((2, 0, 1))
+    batch_image = np.expand_dims(resized_image, axis=0).astype('float32')
+    batch_image = (batch_image / 127.5) - 1.0
+
+    batch_image = torch.from_numpy(batch_image).float().to(device)
+    outputs = model(batch_image)
+
+    pts, pts_score, vmap = deccode_output_score_and_ptss(outputs, 200, 3)
+    start = vmap[:, :, :2]  # (x, y)
+    end = vmap[:, :, 2:]  # (x, y)
+    dist_map = np.sqrt(np.sum((start - end) ** 2, axis=-1))
+
+    junc_list = []
+    segments_list = []
+    for junc, score in zip(pts, pts_score, strict=False):
+        y, x = junc
+        distance = dist_map[y, x]
+        if score > params['score'] and distance > 20.0:
+            junc_list.append([x, y])
+            disp_x_start, disp_y_start, disp_x_end, disp_y_end = vmap[y, x, :]
+            d_arrow = 1.0
+            x_start = x + d_arrow * disp_x_start
+            y_start = y + d_arrow * disp_y_start
+            x_end = x + d_arrow * disp_x_end
+            y_end = y + d_arrow * disp_y_end
+            segments_list.append([x_start, y_start, x_end, y_end])
+
+    segments = np.array(segments_list)
+
+    ####### post processing for squares
+    # 1. get unique lines
+    point = np.array([[0, 0]])
+    point = point[0]
+    start = segments[:, :2]
+    end = segments[:, 2:]
+    diff = start - end
+    a = diff[:, 1]
+    b = -diff[:, 0]
+    c = a * start[:, 0] + b * start[:, 1]
+
+    d = np.abs(a * point[0] + b * point[1] - c) / np.sqrt(a ** 2 + b ** 2 + 1e-10)
+    theta = np.arctan2(diff[:, 0], diff[:, 1]) * 180 / np.pi
+    theta[theta < 0.0] += 180
+    hough = np.concatenate([d[:, None], theta[:, None]], axis=-1)
+
+    d_quant = 1
+    theta_quant = 2
+    hough[:, 0] //= d_quant
+    hough[:, 1] //= theta_quant
+    _, indices, counts = np.unique(hough, axis=0, return_index=True, return_counts=True)
+
+    acc_map = np.zeros([512 // d_quant + 1, 360 // theta_quant + 1], dtype='float32')
+    idx_map = np.zeros([512 // d_quant + 1, 360 // theta_quant + 1], dtype='int32') - 1
+    yx_indices = hough[indices, :].astype('int32')
+    acc_map[yx_indices[:, 0], yx_indices[:, 1]] = counts
+    idx_map[yx_indices[:, 0], yx_indices[:, 1]] = indices
+
+    acc_map_np = acc_map
+    # acc_map = acc_map[None, :, :, None]
+    #
+    # ### fast suppression using tensorflow op
+    # acc_map = tf.constant(acc_map, dtype=tf.float32)
+    # max_acc_map = tf.keras.layers.MaxPool2D(pool_size=(5, 5), strides=1, padding='same')(acc_map)
+    # acc_map = acc_map * tf.cast(tf.math.equal(acc_map, max_acc_map), tf.float32)
+    # flatten_acc_map = tf.reshape(acc_map, [1, -1])
+    # topk_values, topk_indices = tf.math.top_k(flatten_acc_map, k=len(pts))
+    # _, h, w, _ = acc_map.shape
+    # y = tf.expand_dims(topk_indices // w, axis=-1)
+    # x = tf.expand_dims(topk_indices % w, axis=-1)
+    # yx = tf.concat([y, x], axis=-1)
+
+    ### fast suppression using pytorch op
+    acc_map = torch.from_numpy(acc_map_np).unsqueeze(0).unsqueeze(0)
+    _,_, h, w = acc_map.shape
+    max_acc_map = F.max_pool2d(acc_map,kernel_size=5, stride=1, padding=2)
+    acc_map = acc_map * ( (acc_map == max_acc_map).float() )
+    flatten_acc_map = acc_map.reshape([-1, ])
+
+    scores, indices = torch.topk(flatten_acc_map, len(pts), dim=-1, largest=True)
+    yy = torch.div(indices, w, rounding_mode='floor').unsqueeze(-1)
+    xx = torch.fmod(indices, w).unsqueeze(-1)
+    yx = torch.cat((yy, xx), dim=-1)
+
+    yx = yx.detach().cpu().numpy()
+
+    topk_values = scores.detach().cpu().numpy()
+    indices = idx_map[yx[:, 0], yx[:, 1]]
+    basis = 5 // 2
+
+    merged_segments = []
+    for yx_pt, max_indice, value in zip(yx, indices, topk_values, strict=False):
+        y, x = yx_pt
+        if max_indice == -1 or value == 0:
+            continue
+        segment_list = []
+        for y_offset in range(-basis, basis + 1):
+            for x_offset in range(-basis, basis + 1):
+                indice = idx_map[y + y_offset, x + x_offset]
+                cnt = int(acc_map_np[y + y_offset, x + x_offset])
+                if indice != -1:
+                    segment_list.append(segments[indice])
+                if cnt > 1:
+                    check_cnt = 1
+                    current_hough = hough[indice]
+                    for new_indice, new_hough in enumerate(hough):
+                        if (current_hough == new_hough).all() and indice != new_indice:
+                            segment_list.append(segments[new_indice])
+                            check_cnt += 1
+                        if check_cnt == cnt:
+                            break
+        group_segments = np.array(segment_list).reshape([-1, 2])
+        sorted_group_segments = np.sort(group_segments, axis=0)
+        x_min, y_min = sorted_group_segments[0, :]
+        x_max, y_max = sorted_group_segments[-1, :]
+
+        deg = theta[max_indice]
+        if deg >= 90:
+            merged_segments.append([x_min, y_max, x_max, y_min])
+        else:
+            merged_segments.append([x_min, y_min, x_max, y_max])
+
+    # 2. get intersections
+    new_segments = np.array(merged_segments)  # (x1, y1, x2, y2)
+    start = new_segments[:, :2]  # (x1, y1)
+    end = new_segments[:, 2:]  # (x2, y2)
+    new_centers = (start + end) / 2.0
+    diff = start - end
+    dist_segments = np.sqrt(np.sum(diff ** 2, axis=-1))
+
+    # ax + by = c
+    a = diff[:, 1]
+    b = -diff[:, 0]
+    c = a * start[:, 0] + b * start[:, 1]
+    pre_det = a[:, None] * b[None, :]
+    det = pre_det - np.transpose(pre_det)
+
+    pre_inter_y = a[:, None] * c[None, :]
+    inter_y = (pre_inter_y - np.transpose(pre_inter_y)) / (det + 1e-10)
+    pre_inter_x = c[:, None] * b[None, :]
+    inter_x = (pre_inter_x - np.transpose(pre_inter_x)) / (det + 1e-10)
+    inter_pts = np.concatenate([inter_x[:, :, None], inter_y[:, :, None]], axis=-1).astype('int32')
+
+    # 3. get corner information
+    # 3.1 get distance
+    '''
+    dist_segments:
+        | dist(0), dist(1), dist(2), ...|
+    dist_inter_to_segment1:
+        | dist(inter,0), dist(inter,0), dist(inter,0), ... |
+        | dist(inter,1), dist(inter,1), dist(inter,1), ... |
+        ...
+    dist_inter_to_semgnet2:
+        | dist(inter,0), dist(inter,1), dist(inter,2), ... |
+        | dist(inter,0), dist(inter,1), dist(inter,2), ... |
+        ...
+    '''
+
+    dist_inter_to_segment1_start = np.sqrt(
+        np.sum(((inter_pts - start[:, None, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+    dist_inter_to_segment1_end = np.sqrt(
+        np.sum(((inter_pts - end[:, None, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+    dist_inter_to_segment2_start = np.sqrt(
+        np.sum(((inter_pts - start[None, :, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+    dist_inter_to_segment2_end = np.sqrt(
+        np.sum(((inter_pts - end[None, :, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+
+    # sort ascending
+    dist_inter_to_segment1 = np.sort(
+        np.concatenate([dist_inter_to_segment1_start, dist_inter_to_segment1_end], axis=-1),
+        axis=-1)  # [n_batch, n_batch, 2]
+    dist_inter_to_segment2 = np.sort(
+        np.concatenate([dist_inter_to_segment2_start, dist_inter_to_segment2_end], axis=-1),
+        axis=-1)  # [n_batch, n_batch, 2]
+
+    # 3.2 get degree
+    inter_to_start = new_centers[:, None, :] - inter_pts
+    deg_inter_to_start = np.arctan2(inter_to_start[:, :, 1], inter_to_start[:, :, 0]) * 180 / np.pi
+    deg_inter_to_start[deg_inter_to_start < 0.0] += 360
+    inter_to_end = new_centers[None, :, :] - inter_pts
+    deg_inter_to_end = np.arctan2(inter_to_end[:, :, 1], inter_to_end[:, :, 0]) * 180 / np.pi
+    deg_inter_to_end[deg_inter_to_end < 0.0] += 360
+
+    '''
+    B -- G
+    |    |
+    C -- R
+    B : blue / G: green / C: cyan / R: red
+
+    0 -- 1
+    |    |
+    3 -- 2
+    '''
+    # rename variables
+    deg1_map, deg2_map = deg_inter_to_start, deg_inter_to_end
+    # sort deg ascending
+    deg_sort = np.sort(np.concatenate([deg1_map[:, :, None], deg2_map[:, :, None]], axis=-1), axis=-1)
+
+    deg_diff_map = np.abs(deg1_map - deg2_map)
+    # we only consider the smallest degree of intersect
+    deg_diff_map[deg_diff_map > 180] = 360 - deg_diff_map[deg_diff_map > 180]
+
+    # define available degree range
+    deg_range = [60, 120]
+
+    corner_dict = {corner_info: [] for corner_info in range(4)}
+    inter_points = []
+    for i in range(inter_pts.shape[0]):
+        for j in range(i + 1, inter_pts.shape[1]):
+            # i, j > line index, always i < j
+            x, y = inter_pts[i, j, :]
+            deg1, deg2 = deg_sort[i, j, :]
+            deg_diff = deg_diff_map[i, j]
+
+            check_degree = deg_diff > deg_range[0] and deg_diff < deg_range[1]
+
+            outside_ratio = params['outside_ratio']  # over ratio >>> drop it!
+            inside_ratio = params['inside_ratio']  # over ratio >>> drop it!
+            check_distance = ((dist_inter_to_segment1[i, j, 1] >= dist_segments[i] and \
+                               dist_inter_to_segment1[i, j, 0] <= dist_segments[i] * outside_ratio) or \
+                              (dist_inter_to_segment1[i, j, 1] <= dist_segments[i] and \
+                               dist_inter_to_segment1[i, j, 0] <= dist_segments[i] * inside_ratio)) and \
+                             ((dist_inter_to_segment2[i, j, 1] >= dist_segments[j] and \
+                               dist_inter_to_segment2[i, j, 0] <= dist_segments[j] * outside_ratio) or \
+                              (dist_inter_to_segment2[i, j, 1] <= dist_segments[j] and \
+                               dist_inter_to_segment2[i, j, 0] <= dist_segments[j] * inside_ratio))
+
+            if check_degree and check_distance:
+                corner_info = None
+
+                if (deg1 >= 0 and deg1 <= 45 and deg2 >= 45 and deg2 <= 120) or \
+                        (deg2 >= 315 and deg1 >= 45 and deg1 <= 120):
+                    corner_info, color_info = 0, 'blue'
+                elif (deg1 >= 45 and deg1 <= 125 and deg2 >= 125 and deg2 <= 225):
+                    corner_info, color_info = 1, 'green'
+                elif (deg1 >= 125 and deg1 <= 225 and deg2 >= 225 and deg2 <= 315):
+                    corner_info, color_info = 2, 'black'
+                elif (deg1 >= 0 and deg1 <= 45 and deg2 >= 225 and deg2 <= 315) or \
+                        (deg2 >= 315 and deg1 >= 225 and deg1 <= 315):
+                    corner_info, color_info = 3, 'cyan'
+                else:
+                    corner_info, color_info = 4, 'red'  # we don't use it
+                    continue
+
+                corner_dict[corner_info].append([x, y, i, j])
+                inter_points.append([x, y])
+
+    square_list = []
+    connect_list = []
+    segments_list = []
+    for corner0 in corner_dict[0]:
+        for corner1 in corner_dict[1]:
+            connect01 = False
+            for corner0_line in corner0[2:]:
+                if corner0_line in corner1[2:]:
+                    connect01 = True
+                    break
+            if connect01:
+                for corner2 in corner_dict[2]:
+                    connect12 = False
+                    for corner1_line in corner1[2:]:
+                        if corner1_line in corner2[2:]:
+                            connect12 = True
+                            break
+                    if connect12:
+                        for corner3 in corner_dict[3]:
+                            connect23 = False
+                            for corner2_line in corner2[2:]:
+                                if corner2_line in corner3[2:]:
+                                    connect23 = True
+                                    break
+                            if connect23:
+                                for corner3_line in corner3[2:]:
+                                    if corner3_line in corner0[2:]:
+                                        # SQUARE!!!
+                                        '''
+                                        0 -- 1
+                                        |    |
+                                        3 -- 2
+                                        square_list:
+                                            order: 0 > 1 > 2 > 3
+                                            | x0, y0, x1, y1, x2, y2, x3, y3 |
+                                            | x0, y0, x1, y1, x2, y2, x3, y3 |
+                                            ...
+                                        connect_list:
+                                            order: 01 > 12 > 23 > 30
+                                            | line_idx01, line_idx12, line_idx23, line_idx30 |
+                                            | line_idx01, line_idx12, line_idx23, line_idx30 |
+                                            ...
+                                        segments_list:
+                                            order: 0 > 1 > 2 > 3
+                                            | line_idx0_i, line_idx0_j, line_idx1_i, line_idx1_j, line_idx2_i, line_idx2_j, line_idx3_i, line_idx3_j |
+                                            | line_idx0_i, line_idx0_j, line_idx1_i, line_idx1_j, line_idx2_i, line_idx2_j, line_idx3_i, line_idx3_j |
+                                            ...
+                                        '''
+                                        square_list.append(corner0[:2] + corner1[:2] + corner2[:2] + corner3[:2])
+                                        connect_list.append([corner0_line, corner1_line, corner2_line, corner3_line])
+                                        segments_list.append(corner0[2:] + corner1[2:] + corner2[2:] + corner3[2:])
+
+    def check_outside_inside(segments_info, connect_idx):
+        # return 'outside or inside', min distance, cover_param, peri_param
+        if connect_idx == segments_info[0]:
+            check_dist_mat = dist_inter_to_segment1
+        else:
+            check_dist_mat = dist_inter_to_segment2
+
+        i, j = segments_info
+        min_dist, max_dist = check_dist_mat[i, j, :]
+        connect_dist = dist_segments[connect_idx]
+        if max_dist > connect_dist:
+            return 'outside', min_dist, 0, 1
+        else:
+            return 'inside', min_dist, -1, -1
+
+    top_square = None
+
+    try:
+        map_size = input_shape[0] / 2
+        squares = np.array(square_list).reshape([-1, 4, 2])
+        score_array = []
+        connect_array = np.array(connect_list)
+        segments_array = np.array(segments_list).reshape([-1, 4, 2])
+
+        # get degree of corners:
+        squares_rollup = np.roll(squares, 1, axis=1)
+        squares_rolldown = np.roll(squares, -1, axis=1)
+        vec1 = squares_rollup - squares
+        normalized_vec1 = vec1 / (np.linalg.norm(vec1, axis=-1, keepdims=True) + 1e-10)
+        vec2 = squares_rolldown - squares
+        normalized_vec2 = vec2 / (np.linalg.norm(vec2, axis=-1, keepdims=True) + 1e-10)
+        inner_products = np.sum(normalized_vec1 * normalized_vec2, axis=-1)  # [n_squares, 4]
+        squares_degree = np.arccos(inner_products) * 180 / np.pi  # [n_squares, 4]
+
+        # get square score
+        overlap_scores = []
+        degree_scores = []
+        length_scores = []
+
+        for connects, segments, square, degree in zip(connect_array, segments_array, squares, squares_degree, strict=False):
+            '''
+            0 -- 1
+            |    |
+            3 -- 2
+
+            # segments: [4, 2]
+            # connects: [4]
+            '''
+
+            ###################################### OVERLAP SCORES
+            cover = 0
+            perimeter = 0
+            # check 0 > 1 > 2 > 3
+            square_length = []
+
+            for start_idx in range(4):
+                end_idx = (start_idx + 1) % 4
+
+                connect_idx = connects[start_idx]  # segment idx of segment01
+                start_segments = segments[start_idx]
+                end_segments = segments[end_idx]
+
+                start_point = square[start_idx]
+                end_point = square[end_idx]
+
+                # check whether outside or inside
+                start_position, start_min, start_cover_param, start_peri_param = check_outside_inside(start_segments,
+                                                                                                      connect_idx)
+                end_position, end_min, end_cover_param, end_peri_param = check_outside_inside(end_segments, connect_idx)
+
+                cover += dist_segments[connect_idx] + start_cover_param * start_min + end_cover_param * end_min
+                perimeter += dist_segments[connect_idx] + start_peri_param * start_min + end_peri_param * end_min
+
+                square_length.append(
+                    dist_segments[connect_idx] + start_peri_param * start_min + end_peri_param * end_min)
+
+            overlap_scores.append(cover / perimeter)
+            ######################################
+            ###################################### DEGREE SCORES
+            '''
+            deg0 vs deg2
+            deg1 vs deg3
+            '''
+            deg0, deg1, deg2, deg3 = degree
+            deg_ratio1 = deg0 / deg2
+            if deg_ratio1 > 1.0:
+                deg_ratio1 = 1 / deg_ratio1
+            deg_ratio2 = deg1 / deg3
+            if deg_ratio2 > 1.0:
+                deg_ratio2 = 1 / deg_ratio2
+            degree_scores.append((deg_ratio1 + deg_ratio2) / 2)
+            ######################################
+            ###################################### LENGTH SCORES
+            '''
+            len0 vs len2
+            len1 vs len3
+            '''
+            len0, len1, len2, len3 = square_length
+            len_ratio1 = len0 / len2 if len2 > len0 else len2 / len0
+            len_ratio2 = len1 / len3 if len3 > len1 else len3 / len1
+            length_scores.append((len_ratio1 + len_ratio2) / 2)
+
+            ######################################
+
+        overlap_scores = np.array(overlap_scores)
+        overlap_scores /= np.max(overlap_scores)
+
+        degree_scores = np.array(degree_scores)
+        # degree_scores /= np.max(degree_scores)
+
+        length_scores = np.array(length_scores)
+
+        ###################################### AREA SCORES
+        area_scores = np.reshape(squares, [-1, 4, 2])
+        area_x = area_scores[:, :, 0]
+        area_y = area_scores[:, :, 1]
+        correction = area_x[:, -1] * area_y[:, 0] - area_y[:, -1] * area_x[:, 0]
+        area_scores = np.sum(area_x[:, :-1] * area_y[:, 1:], axis=-1) - np.sum(area_y[:, :-1] * area_x[:, 1:], axis=-1)
+        area_scores = 0.5 * np.abs(area_scores + correction)
+        area_scores /= (map_size * map_size)  # np.max(area_scores)
+        ######################################
+
+        ###################################### CENTER SCORES
+        centers = np.array([[256 // 2, 256 // 2]], dtype='float32')  # [1, 2]
+        # squares: [n, 4, 2]
+        square_centers = np.mean(squares, axis=1)  # [n, 2]
+        center2center = np.sqrt(np.sum((centers - square_centers) ** 2))
+        center_scores = center2center / (map_size / np.sqrt(2.0))
+
+        '''
+        score_w = [overlap, degree, area, center, length]
+        '''
+        score_w = [0.0, 1.0, 10.0, 0.5, 1.0]
+        score_array = params['w_overlap'] * overlap_scores \
+                      + params['w_degree'] * degree_scores \
+                      + params['w_area'] * area_scores \
+                      - params['w_center'] * center_scores \
+                      + params['w_length'] * length_scores
+
+        best_square = []
+
+        sorted_idx = np.argsort(score_array)[::-1]
+        score_array = score_array[sorted_idx]
+        squares = squares[sorted_idx]
+
+    except Exception:
+        pass
+
+    '''return list
+    merged_lines, squares, scores
+    '''
+
+    try:
+        new_segments[:, 0] = new_segments[:, 0] * 2 / input_shape[1] * original_shape[1]
+        new_segments[:, 1] = new_segments[:, 1] * 2 / input_shape[0] * original_shape[0]
+        new_segments[:, 2] = new_segments[:, 2] * 2 / input_shape[1] * original_shape[1]
+        new_segments[:, 3] = new_segments[:, 3] * 2 / input_shape[0] * original_shape[0]
+    except Exception:
+        new_segments = []
+
+    try:
+        squares[:, :, 0] = squares[:, :, 0] * 2 / input_shape[1] * original_shape[1]
+        squares[:, :, 1] = squares[:, :, 1] * 2 / input_shape[0] * original_shape[0]
+    except Exception:
+        squares = []
+        score_array = []
+
+    try:
+        inter_points = np.array(inter_points)
+        inter_points[:, 0] = inter_points[:, 0] * 2 / input_shape[1] * original_shape[1]
+        inter_points[:, 1] = inter_points[:, 1] * 2 / input_shape[0] * original_shape[0]
+    except Exception:
+        inter_points = []
+
+    return new_segments, squares, score_array, inter_points

invokeai/backend/image_util/normal_bae/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Caroline Chan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

invokeai/backend/image_util/normal_bae/__init__.py

@@ -0,0 +1,94 @@
+# Adapted from https://github.com/huggingface/controlnet_aux
+
+import pathlib
+import types
+
+import cv2
+import huggingface_hub
+import numpy as np
+import torch
+import torchvision.transforms as transforms
+from einops import rearrange
+from PIL import Image
+
+from invokeai.backend.image_util.normal_bae.nets.NNET import NNET
+from invokeai.backend.image_util.util import np_to_pil, pil_to_np, resize_to_multiple
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+class NormalMapDetector:
+    """Simple wrapper around the Normal BAE model for normal map generation."""
+
+    hf_repo_id = "lllyasviel/Annotators"
+    hf_filename = "scannet.pt"
+
+    @classmethod
+    def get_model_url(cls) -> str:
+        """Get the URL to download the model from the Hugging Face Hub."""
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path) -> NNET:
+        """Load the model from a file."""
+
+        args = types.SimpleNamespace()
+        args.mode = "client"
+        args.architecture = "BN"
+        args.pretrained = "scannet"
+        args.sampling_ratio = 0.4
+        args.importance_ratio = 0.7
+
+        model = NNET(args)
+
+        ckpt = torch.load(model_path, map_location="cpu")["model"]
+        load_dict = {}
+        for k, v in ckpt.items():
+            if k.startswith("module."):
+                k_ = k.replace("module.", "")
+                load_dict[k_] = v
+            else:
+                load_dict[k] = v
+
+        model.load_state_dict(load_dict)
+        model.eval()
+
+        return model
+
+    def __init__(self, model: NNET) -> None:
+        self.model = model
+        self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, image: Image.Image):
+        """Processes an image and returns the detected normal map."""
+
+        device = get_effective_device(self.model)
+        np_image = pil_to_np(image)
+
+        height, width, _channels = np_image.shape
+
+        # The model requires the image to be a multiple of 8
+        np_image = resize_to_multiple(np_image, 8)
+
+        image_normal = np_image
+
+        with torch.no_grad():
+            image_normal = torch.from_numpy(image_normal).float().to(device)
+            image_normal = image_normal / 255.0
+            image_normal = rearrange(image_normal, "h w c -> 1 c h w")
+            image_normal = self.norm(image_normal)
+
+            normal = self.model(image_normal)
+            normal = normal[0][-1][:, :3]
+            normal = ((normal + 1) * 0.5).clip(0, 1)
+
+            normal = rearrange(normal[0], "c h w -> h w c").cpu().numpy()
+            normal_image = (normal * 255.0).clip(0, 255).astype(np.uint8)
+
+        # Back to the original size
+        output_image = cv2.resize(normal_image, (width, height), interpolation=cv2.INTER_LINEAR)
+
+        return np_to_pil(output_image)

invokeai/backend/image_util/normal_bae/nets/NNET.py

@@ -0,0 +1,22 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .submodules.encoder import Encoder
+from .submodules.decoder import Decoder
+
+
+class NNET(nn.Module):
+    def __init__(self, args):
+        super(NNET, self).__init__()
+        self.encoder = Encoder()
+        self.decoder = Decoder(args)
+
+    def get_1x_lr_params(self):  # lr/10 learning rate
+        return self.encoder.parameters()
+
+    def get_10x_lr_params(self):  # lr learning rate
+        return self.decoder.parameters()
+
+    def forward(self, img, **kwargs):
+        return self.decoder(self.encoder(img), **kwargs)

invokeai/backend/image_util/normal_bae/nets/baseline.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .submodules.submodules import UpSampleBN, norm_normalize
+
+
+# This is the baseline encoder-decoder we used in the ablation study
+class NNET(nn.Module):
+    def __init__(self, args=None):
+        super(NNET, self).__init__()
+        self.encoder = Encoder()
+        self.decoder = Decoder(num_classes=4)
+
+    def forward(self, x, **kwargs):
+        out = self.decoder(self.encoder(x), **kwargs)
+
+        # Bilinearly upsample the output to match the input resolution
+        up_out = F.interpolate(out, size=[x.size(2), x.size(3)], mode='bilinear', align_corners=False)
+        
+        # L2-normalize the first three channels / ensure positive value for concentration parameters (kappa)
+        up_out = norm_normalize(up_out) 
+        return up_out
+
+    def get_1x_lr_params(self):  # lr/10 learning rate
+        return self.encoder.parameters()
+
+    def get_10x_lr_params(self):  # lr learning rate
+        modules = [self.decoder]
+        for m in modules:
+            yield from m.parameters()
+
+
+# Encoder
+class Encoder(nn.Module):
+    def __init__(self):
+        super(Encoder, self).__init__()
+
+        basemodel_name = 'tf_efficientnet_b5_ap'
+        basemodel = torch.hub.load('rwightman/gen-efficientnet-pytorch', basemodel_name, pretrained=True)
+
+        # Remove last layer
+        basemodel.global_pool = nn.Identity()
+        basemodel.classifier = nn.Identity()
+
+        self.original_model = basemodel
+
+    def forward(self, x):
+        features = [x]
+        for k, v in self.original_model._modules.items():
+            if (k == 'blocks'):
+                for ki, vi in v._modules.items():
+                    features.append(vi(features[-1]))
+            else:
+                features.append(v(features[-1]))
+        return features
+
+
+# Decoder (no pixel-wise MLP, no uncertainty-guided sampling)
+class Decoder(nn.Module):
+    def __init__(self, num_classes=4):
+        super(Decoder, self).__init__()
+        self.conv2 = nn.Conv2d(2048, 2048, kernel_size=1, stride=1, padding=0)
+        self.up1 = UpSampleBN(skip_input=2048 + 176, output_features=1024)
+        self.up2 = UpSampleBN(skip_input=1024 + 64, output_features=512)
+        self.up3 = UpSampleBN(skip_input=512 + 40, output_features=256)
+        self.up4 = UpSampleBN(skip_input=256 + 24, output_features=128)
+        self.conv3 = nn.Conv2d(128, num_classes, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, features):
+        x_block0, x_block1, x_block2, x_block3, x_block4 = features[4], features[5], features[6], features[8], features[11]
+        x_d0 = self.conv2(x_block4)
+        x_d1 = self.up1(x_d0, x_block3)
+        x_d2 = self.up2(x_d1, x_block2)
+        x_d3 = self.up3(x_d2, x_block1)
+        x_d4 = self.up4(x_d3, x_block0)
+        out = self.conv3(x_d4)
+        return out
+
+
+if __name__ == '__main__':
+    model = Baseline()
+    x = torch.rand(2, 3, 480, 640)
+    out = model(x)
+    print(out.shape)

invokeai/backend/image_util/normal_bae/nets/submodules/decoder.py

@@ -0,0 +1,202 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .submodules import UpSampleBN, UpSampleGN, norm_normalize, sample_points
+
+
+class Decoder(nn.Module):
+    def __init__(self, args):
+        super(Decoder, self).__init__()
+
+        # hyper-parameter for sampling
+        self.sampling_ratio = args.sampling_ratio
+        self.importance_ratio = args.importance_ratio
+
+        # feature-map
+        self.conv2 = nn.Conv2d(2048, 2048, kernel_size=1, stride=1, padding=0)
+        if args.architecture == 'BN':
+            self.up1 = UpSampleBN(skip_input=2048 + 176, output_features=1024)
+            self.up2 = UpSampleBN(skip_input=1024 + 64, output_features=512)
+            self.up3 = UpSampleBN(skip_input=512 + 40, output_features=256)
+            self.up4 = UpSampleBN(skip_input=256 + 24, output_features=128)
+
+        elif args.architecture == 'GN':
+            self.up1 = UpSampleGN(skip_input=2048 + 176, output_features=1024)
+            self.up2 = UpSampleGN(skip_input=1024 + 64, output_features=512)
+            self.up3 = UpSampleGN(skip_input=512 + 40, output_features=256)
+            self.up4 = UpSampleGN(skip_input=256 + 24, output_features=128)
+
+        else:
+            raise Exception('invalid architecture')
+
+        # produces 1/8 res output
+        self.out_conv_res8 = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
+
+        # produces 1/4 res output
+        self.out_conv_res4 = nn.Sequential(
+            nn.Conv1d(512 + 4, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 4, kernel_size=1),
+        )
+
+        # produces 1/2 res output
+        self.out_conv_res2 = nn.Sequential(
+            nn.Conv1d(256 + 4, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 4, kernel_size=1),
+        )
+
+        # produces 1/1 res output
+        self.out_conv_res1 = nn.Sequential(
+            nn.Conv1d(128 + 4, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 4, kernel_size=1),
+        )
+
+    def forward(self, features, gt_norm_mask=None, mode='test'):
+        x_block0, x_block1, x_block2, x_block3, x_block4 = features[4], features[5], features[6], features[8], features[11]
+
+        # generate feature-map
+
+        x_d0 = self.conv2(x_block4)                     # x_d0 : [2, 2048, 15, 20]      1/32 res
+        x_d1 = self.up1(x_d0, x_block3)                 # x_d1 : [2, 1024, 30, 40]      1/16 res
+        x_d2 = self.up2(x_d1, x_block2)                 # x_d2 : [2, 512, 60, 80]       1/8 res
+        x_d3 = self.up3(x_d2, x_block1)                 # x_d3: [2, 256, 120, 160]      1/4 res
+        x_d4 = self.up4(x_d3, x_block0)                 # x_d4: [2, 128, 240, 320]      1/2 res
+
+        # 1/8 res output
+        out_res8 = self.out_conv_res8(x_d2)             # out_res8: [2, 4, 60, 80]      1/8 res output
+        out_res8 = norm_normalize(out_res8)             # out_res8: [2, 4, 60, 80]      1/8 res output
+
+        ################################################################################################################
+        # out_res4
+        ################################################################################################################
+
+        if mode == 'train':
+            # upsampling ... out_res8: [2, 4, 60, 80] -> out_res8_res4: [2, 4, 120, 160]
+            out_res8_res4 = F.interpolate(out_res8, scale_factor=2, mode='bilinear', align_corners=True)
+            B, _, H, W = out_res8_res4.shape
+
+            # samples: [B, 1, N, 2]
+            point_coords_res4, rows_int, cols_int = sample_points(out_res8_res4.detach(), gt_norm_mask,
+                                                                  sampling_ratio=self.sampling_ratio,
+                                                                  beta=self.importance_ratio)
+
+            # output (needed for evaluation / visualization)
+            out_res4 = out_res8_res4
+
+            # grid_sample feature-map
+            feat_res4 = F.grid_sample(x_d2, point_coords_res4, mode='bilinear', align_corners=True)  # (B, 512, 1, N)
+            init_pred = F.grid_sample(out_res8, point_coords_res4, mode='bilinear', align_corners=True)  # (B, 4, 1, N)
+            feat_res4 = torch.cat([feat_res4, init_pred], dim=1)  # (B, 512+4, 1, N)
+
+            # prediction (needed to compute loss)
+            samples_pred_res4 = self.out_conv_res4(feat_res4[:, :, 0, :])  # (B, 4, N)
+            samples_pred_res4 = norm_normalize(samples_pred_res4)  # (B, 4, N) - normalized
+
+            for i in range(B):
+                out_res4[i, :, rows_int[i, :], cols_int[i, :]] = samples_pred_res4[i, :, :]
+
+        else:
+            # grid_sample feature-map
+            feat_map = F.interpolate(x_d2, scale_factor=2, mode='bilinear', align_corners=True)
+            init_pred = F.interpolate(out_res8, scale_factor=2, mode='bilinear', align_corners=True)
+            feat_map = torch.cat([feat_map, init_pred], dim=1)  # (B, 512+4, H, W)
+            B, _, H, W = feat_map.shape
+
+            # try all pixels
+            out_res4 = self.out_conv_res4(feat_map.view(B, 512 + 4, -1))  # (B, 4, N)
+            out_res4 = norm_normalize(out_res4)  # (B, 4, N) - normalized
+            out_res4 = out_res4.view(B, 4, H, W)
+            samples_pred_res4 = point_coords_res4 = None
+
+        ################################################################################################################
+        # out_res2
+        ################################################################################################################
+
+        if mode == 'train':
+
+            # upsampling ... out_res4: [2, 4, 120, 160] -> out_res4_res2: [2, 4, 240, 320]
+            out_res4_res2 = F.interpolate(out_res4, scale_factor=2, mode='bilinear', align_corners=True)
+            B, _, H, W = out_res4_res2.shape
+
+            # samples: [B, 1, N, 2]
+            point_coords_res2, rows_int, cols_int = sample_points(out_res4_res2.detach(), gt_norm_mask,
+                                                                  sampling_ratio=self.sampling_ratio,
+                                                                  beta=self.importance_ratio)
+
+            # output (needed for evaluation / visualization)
+            out_res2 = out_res4_res2
+
+            # grid_sample feature-map
+            feat_res2 = F.grid_sample(x_d3, point_coords_res2, mode='bilinear', align_corners=True)  # (B, 256, 1, N)
+            init_pred = F.grid_sample(out_res4, point_coords_res2, mode='bilinear', align_corners=True)  # (B, 4, 1, N)
+            feat_res2 = torch.cat([feat_res2, init_pred], dim=1)  # (B, 256+4, 1, N)
+
+            # prediction (needed to compute loss)
+            samples_pred_res2 = self.out_conv_res2(feat_res2[:, :, 0, :])  # (B, 4, N)
+            samples_pred_res2 = norm_normalize(samples_pred_res2)  # (B, 4, N) - normalized
+
+            for i in range(B):
+                out_res2[i, :, rows_int[i, :], cols_int[i, :]] = samples_pred_res2[i, :, :]
+
+        else:
+            # grid_sample feature-map
+            feat_map = F.interpolate(x_d3, scale_factor=2, mode='bilinear', align_corners=True)
+            init_pred = F.interpolate(out_res4, scale_factor=2, mode='bilinear', align_corners=True)
+            feat_map = torch.cat([feat_map, init_pred], dim=1)  # (B, 512+4, H, W)
+            B, _, H, W = feat_map.shape
+
+            out_res2 = self.out_conv_res2(feat_map.view(B, 256 + 4, -1))  # (B, 4, N)
+            out_res2 = norm_normalize(out_res2)  # (B, 4, N) - normalized
+            out_res2 = out_res2.view(B, 4, H, W)
+            samples_pred_res2 = point_coords_res2 = None
+
+        ################################################################################################################
+        # out_res1
+        ################################################################################################################
+
+        if mode == 'train':
+            # upsampling ... out_res4: [2, 4, 120, 160] -> out_res4_res2: [2, 4, 240, 320]
+            out_res2_res1 = F.interpolate(out_res2, scale_factor=2, mode='bilinear', align_corners=True)
+            B, _, H, W = out_res2_res1.shape
+
+            # samples: [B, 1, N, 2]
+            point_coords_res1, rows_int, cols_int = sample_points(out_res2_res1.detach(), gt_norm_mask,
+                                                                  sampling_ratio=self.sampling_ratio,
+                                                                  beta=self.importance_ratio)
+
+            # output (needed for evaluation / visualization)
+            out_res1 = out_res2_res1
+
+            # grid_sample feature-map
+            feat_res1 = F.grid_sample(x_d4, point_coords_res1, mode='bilinear', align_corners=True)  # (B, 128, 1, N)
+            init_pred = F.grid_sample(out_res2, point_coords_res1, mode='bilinear', align_corners=True)  # (B, 4, 1, N)
+            feat_res1 = torch.cat([feat_res1, init_pred], dim=1)  # (B, 128+4, 1, N)
+
+            # prediction (needed to compute loss)
+            samples_pred_res1 = self.out_conv_res1(feat_res1[:, :, 0, :])  # (B, 4, N)
+            samples_pred_res1 = norm_normalize(samples_pred_res1)  # (B, 4, N) - normalized
+
+            for i in range(B):
+                out_res1[i, :, rows_int[i, :], cols_int[i, :]] = samples_pred_res1[i, :, :]
+
+        else:
+            # grid_sample feature-map
+            feat_map = F.interpolate(x_d4, scale_factor=2, mode='bilinear', align_corners=True)
+            init_pred = F.interpolate(out_res2, scale_factor=2, mode='bilinear', align_corners=True)
+            feat_map = torch.cat([feat_map, init_pred], dim=1)  # (B, 512+4, H, W)
+            B, _, H, W = feat_map.shape
+
+            out_res1 = self.out_conv_res1(feat_map.view(B, 128 + 4, -1))  # (B, 4, N)
+            out_res1 = norm_normalize(out_res1)  # (B, 4, N) - normalized
+            out_res1 = out_res1.view(B, 4, H, W)
+            samples_pred_res1 = point_coords_res1 = None
+
+        return [out_res8, out_res4, out_res2, out_res1], \
+               [out_res8, samples_pred_res4, samples_pred_res2, samples_pred_res1], \
+               [None, point_coords_res4, point_coords_res2, point_coords_res1]
+

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/.gitignore

@@ -0,0 +1,109 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# pytorch stuff
+*.pth
+*.onnx
+*.pb
+
+trained_models/
+.fuse_hidden*

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/BENCHMARK.md

@@ -0,0 +1,555 @@
+# Model Performance Benchmarks
+
+All benchmarks run as per:
+
+```
+python onnx_export.py --model mobilenetv3_100 ./mobilenetv3_100.onnx
+python onnx_optimize.py ./mobilenetv3_100.onnx --output mobilenetv3_100-opt.onnx
+python onnx_to_caffe.py ./mobilenetv3_100.onnx --c2-prefix mobilenetv3
+python onnx_to_caffe.py ./mobilenetv3_100-opt.onnx --c2-prefix mobilenetv3-opt
+python caffe2_benchmark.py --c2-init ./mobilenetv3.init.pb --c2-predict ./mobilenetv3.predict.pb
+python caffe2_benchmark.py --c2-init ./mobilenetv3-opt.init.pb --c2-predict ./mobilenetv3-opt.predict.pb
+```
+
+## EfficientNet-B0
+
+### Unoptimized
+```
+Main run finished. Milliseconds per iter: 49.2862. Iters per second: 20.2897
+Time per operator type:
+        29.7378 ms.    60.5145%. Conv
+        12.1785 ms.    24.7824%. Sigmoid
+        3.62811 ms.    7.38297%. SpatialBN
+        2.98444 ms.    6.07314%. Mul
+       0.326902 ms.   0.665225%. AveragePool
+       0.197317 ms.   0.401528%. FC
+      0.0852877 ms.   0.173555%. Add
+      0.0032607 ms. 0.00663532%. Squeeze
+        49.1416 ms in Total
+FLOP per operator type:
+        0.76907 GFLOP.    95.2696%. Conv
+      0.0269508 GFLOP.    3.33857%. SpatialBN
+     0.00846444 GFLOP.    1.04855%. Mul
+       0.002561 GFLOP.   0.317248%. FC
+    0.000210112 GFLOP.  0.0260279%. Add
+       0.807256 GFLOP in Total
+Feature Memory Read per operator type:
+        58.5253 MB.    43.0891%. Mul
+        43.2015 MB.     31.807%. Conv
+        27.2869 MB.    20.0899%. SpatialBN
+        5.12912 MB.    3.77631%. FC
+         1.6809 MB.    1.23756%. Add
+        135.824 MB in Total
+Feature Memory Written per operator type:
+        33.8578 MB.    38.1965%. Mul
+        26.9881 MB.    30.4465%. Conv
+        26.9508 MB.    30.4044%. SpatialBN
+       0.840448 MB.   0.948147%. Add
+          0.004 MB. 0.00451258%. FC
+        88.6412 MB in Total
+Parameter Memory per operator type:
+        15.8248 MB.    74.9391%. Conv
+          5.124 MB.     24.265%. FC
+       0.168064 MB.   0.795877%. SpatialBN
+              0 MB.          0%. Add
+              0 MB.          0%. Mul
+        21.1168 MB in Total
+```
+### Optimized
+```
+Main run finished. Milliseconds per iter: 46.0838. Iters per second: 21.6996
+Time per operator type:
+         29.776 ms.     65.002%. Conv
+        12.2803 ms.    26.8084%. Sigmoid
+        3.15073 ms.    6.87815%. Mul
+       0.328651 ms.   0.717456%. AveragePool
+       0.186237 ms.   0.406563%. FC
+      0.0832429 ms.   0.181722%. Add
+      0.0026184 ms. 0.00571606%. Squeeze
+        45.8078 ms in Total
+FLOP per operator type:
+        0.76907 GFLOP.    98.5601%. Conv
+     0.00846444 GFLOP.    1.08476%. Mul
+       0.002561 GFLOP.   0.328205%. FC
+    0.000210112 GFLOP.  0.0269269%. Add
+       0.780305 GFLOP in Total
+Feature Memory Read per operator type:
+        58.5253 MB.    53.8803%. Mul
+        43.2855 MB.    39.8501%. Conv
+        5.12912 MB.    4.72204%. FC
+         1.6809 MB.    1.54749%. Add
+        108.621 MB in Total
+Feature Memory Written per operator type:
+        33.8578 MB.    54.8834%. Mul
+        26.9881 MB.    43.7477%. Conv
+       0.840448 MB.    1.36237%. Add
+          0.004 MB. 0.00648399%. FC
+        61.6904 MB in Total
+Parameter Memory per operator type:
+        15.8248 MB.    75.5403%. Conv
+          5.124 MB.    24.4597%. FC
+              0 MB.          0%. Add
+              0 MB.          0%. Mul
+        20.9488 MB in Total
+```
+
+## EfficientNet-B1
+### Optimized
+```
+Main run finished. Milliseconds per iter: 71.8102. Iters per second: 13.9256
+Time per operator type:
+        45.7915 ms.    66.3206%. Conv
+        17.8718 ms.    25.8841%. Sigmoid
+        4.44132 ms.    6.43244%. Mul
+        0.51001 ms.   0.738658%. AveragePool
+       0.233283 ms.   0.337868%. Add
+       0.194986 ms.   0.282402%. FC
+     0.00268255 ms. 0.00388519%. Squeeze
+        69.0456 ms in Total
+FLOP per operator type:
+        1.37105 GFLOP.    98.7673%. Conv
+      0.0138759 GFLOP.    0.99959%. Mul
+       0.002561 GFLOP.   0.184489%. FC
+    0.000674432 GFLOP.  0.0485847%. Add
+        1.38816 GFLOP in Total
+Feature Memory Read per operator type:
+         94.624 MB.    54.0789%. Mul
+        69.8255 MB.    39.9062%. Conv
+        5.39546 MB.    3.08357%. Add
+        5.12912 MB.    2.93136%. FC
+        174.974 MB in Total
+Feature Memory Written per operator type:
+        55.5035 MB.     54.555%. Mul
+        43.5333 MB.    42.7894%. Conv
+        2.69773 MB.    2.65163%. Add
+          0.004 MB. 0.00393165%. FC
+        101.739 MB in Total
+Parameter Memory per operator type:
+        25.7479 MB.    83.4024%. Conv
+          5.124 MB.    16.5976%. FC
+              0 MB.          0%. Add
+              0 MB.          0%. Mul
+        30.8719 MB in Total
+```
+
+## EfficientNet-B2
+### Optimized
+```
+Main run finished. Milliseconds per iter: 92.28. Iters per second: 10.8366
+Time per operator type:
+        61.4627 ms.    67.5845%. Conv
+        22.7458 ms.    25.0113%. Sigmoid
+        5.59931 ms.    6.15701%. Mul
+       0.642567 ms.   0.706568%. AveragePool
+       0.272795 ms.   0.299965%. Add
+       0.216178 ms.   0.237709%. FC
+     0.00268895 ms. 0.00295677%. Squeeze
+         90.942 ms in Total
+FLOP per operator type:
+        1.98431 GFLOP.    98.9343%. Conv
+      0.0177039 GFLOP.   0.882686%. Mul
+       0.002817 GFLOP.   0.140451%. FC
+    0.000853984 GFLOP.  0.0425782%. Add
+        2.00568 GFLOP in Total
+Feature Memory Read per operator type:
+        120.609 MB.    54.9637%. Mul
+        86.3512 MB.    39.3519%. Conv
+        6.83187 MB.    3.11341%. Add
+        5.64163 MB.      2.571%. FC
+        219.433 MB in Total
+Feature Memory Written per operator type:
+        70.8155 MB.    54.6573%. Mul
+        55.3273 MB.    42.7031%. Conv
+        3.41594 MB.    2.63651%. Add
+          0.004 MB. 0.00308731%. FC
+        129.563 MB in Total
+Parameter Memory per operator type:
+        30.4721 MB.    84.3913%. Conv
+          5.636 MB.    15.6087%. FC
+              0 MB.          0%. Add
+              0 MB.          0%. Mul
+        36.1081 MB in Total
+```
+
+## MixNet-M
+### Optimized
+```
+Main run finished. Milliseconds per iter: 63.1122. Iters per second: 15.8448
+Time per operator type:
+        48.1139 ms.    75.2052%. Conv
+         7.1341 ms.    11.1511%. Sigmoid
+        2.63706 ms.    4.12189%. SpatialBN
+        1.73186 ms.    2.70701%. Mul
+        1.38707 ms.    2.16809%. Split
+        1.29322 ms.    2.02139%. Concat
+        1.00093 ms.    1.56452%. Relu
+       0.235309 ms.   0.367803%. Add
+       0.221579 ms.   0.346343%. FC
+       0.219315 ms.   0.342803%. AveragePool
+     0.00250145 ms. 0.00390993%. Squeeze
+        63.9768 ms in Total
+FLOP per operator type:
+       0.675273 GFLOP.    95.5827%. Conv
+      0.0221072 GFLOP.    3.12921%. SpatialBN
+     0.00538445 GFLOP.   0.762152%. Mul
+       0.003073 GFLOP.   0.434973%. FC
+    0.000642488 GFLOP.  0.0909421%. Add
+              0 GFLOP.          0%. Concat
+              0 GFLOP.          0%. Relu
+        0.70648 GFLOP in Total
+Feature Memory Read per operator type:
+        46.8424 MB.     30.502%. Conv
+        36.8626 MB.    24.0036%. Mul
+        22.3152 MB.    14.5309%. SpatialBN
+        22.1074 MB.    14.3955%. Concat
+        14.1496 MB.    9.21372%. Relu
+        6.15414 MB.    4.00735%. FC
+         5.1399 MB.    3.34692%. Add
+        153.571 MB in Total
+Feature Memory Written per operator type:
+        32.7672 MB.    28.4331%. Conv
+        22.1072 MB.    19.1831%. Concat
+        22.1072 MB.    19.1831%. SpatialBN
+        21.5378 MB.     18.689%. Mul
+        14.1496 MB.    12.2781%. Relu
+        2.56995 MB.    2.23003%. Add
+          0.004 MB. 0.00347092%. FC
+        115.243 MB in Total
+Parameter Memory per operator type:
+        13.7059 MB.     68.674%. Conv
+          6.148 MB.    30.8049%. FC
+          0.104 MB.   0.521097%. SpatialBN
+              0 MB.          0%. Add
+              0 MB.          0%. Concat
+              0 MB.          0%. Mul
+              0 MB.          0%. Relu
+        19.9579 MB in Total
+```
+
+## TF MobileNet-V3 Large 1.0
+
+### Optimized
+```
+Main run finished. Milliseconds per iter: 22.0495. Iters per second: 45.3525
+Time per operator type:
+         17.437 ms.    80.0087%. Conv
+        1.27662 ms.     5.8577%. Add
+        1.12759 ms.    5.17387%. Div
+       0.701155 ms.    3.21721%. Mul
+       0.562654 ms.    2.58171%. Relu
+       0.431144 ms.    1.97828%. Clip
+       0.156902 ms.   0.719936%. FC
+      0.0996858 ms.   0.457402%. AveragePool
+     0.00112455 ms. 0.00515993%. Flatten
+        21.7939 ms in Total
+FLOP per operator type:
+        0.43062 GFLOP.    98.1484%. Conv
+       0.002561 GFLOP.   0.583713%. FC
+     0.00210867 GFLOP.   0.480616%. Mul
+     0.00193868 GFLOP.   0.441871%. Add
+     0.00151532 GFLOP.   0.345377%. Div
+              0 GFLOP.          0%. Relu
+       0.438743 GFLOP in Total
+Feature Memory Read per operator type:
+        34.7967 MB.    43.9391%. Conv
+         14.496 MB.    18.3046%. Mul
+        9.44828 MB.    11.9307%. Add
+        9.26157 MB.    11.6949%. Relu
+         6.0614 MB.    7.65395%. Div
+        5.12912 MB.    6.47673%. FC
+         79.193 MB in Total
+Feature Memory Written per operator type:
+        17.6247 MB.    35.8656%. Conv
+        9.26157 MB.     18.847%. Relu
+        8.43469 MB.    17.1643%. Mul
+        7.75472 MB.    15.7806%. Add
+        6.06128 MB.    12.3345%. Div
+          0.004 MB. 0.00813985%. FC
+        49.1409 MB in Total
+Parameter Memory per operator type:
+        16.6851 MB.    76.5052%. Conv
+          5.124 MB.    23.4948%. FC
+              0 MB.          0%. Add
+              0 MB.          0%. Div
+              0 MB.          0%. Mul
+              0 MB.          0%. Relu
+        21.8091 MB in Total
+```
+
+## MobileNet-V3 (RW)
+
+### Unoptimized
+```
+Main run finished. Milliseconds per iter: 24.8316. Iters per second: 40.2712
+Time per operator type:
+        15.9266 ms.    69.2624%. Conv
+        2.36551 ms.    10.2873%. SpatialBN
+        1.39102 ms.    6.04936%. Add
+        1.30327 ms.    5.66773%. Div
+       0.737014 ms.    3.20517%. Mul
+       0.639697 ms.    2.78195%. Relu
+       0.375681 ms.    1.63378%. Clip
+       0.153126 ms.   0.665921%. FC
+      0.0993787 ms.   0.432184%. AveragePool
+      0.0032632 ms.  0.0141912%. Squeeze
+        22.9946 ms in Total
+FLOP per operator type:
+       0.430616 GFLOP.    94.4041%. Conv
+      0.0175992 GFLOP.    3.85829%. SpatialBN
+       0.002561 GFLOP.   0.561449%. FC
+     0.00210961 GFLOP.    0.46249%. Mul
+     0.00173891 GFLOP.   0.381223%. Add
+     0.00151626 GFLOP.    0.33241%. Div
+              0 GFLOP.          0%. Relu
+       0.456141 GFLOP in Total
+Feature Memory Read per operator type:
+        34.7354 MB.    36.4363%. Conv
+        17.7944 MB.    18.6658%. SpatialBN
+        14.5035 MB.    15.2137%. Mul
+        9.25778 MB.    9.71113%. Relu
+        7.84641 MB.    8.23064%. Add
+        6.06516 MB.    6.36216%. Div
+        5.12912 MB.    5.38029%. FC
+        95.3317 MB in Total
+Feature Memory Written per operator type:
+        17.6246 MB.    26.7264%. Conv
+        17.5992 MB.    26.6878%. SpatialBN
+        9.25778 MB.    14.0387%. Relu
+        8.43843 MB.    12.7962%. Mul
+        6.95565 MB.    10.5477%. Add
+        6.06502 MB.    9.19713%. Div
+          0.004 MB. 0.00606568%. FC
+        65.9447 MB in Total
+Parameter Memory per operator type:
+        16.6778 MB.    76.1564%. Conv
+          5.124 MB.    23.3979%. FC
+         0.0976 MB.   0.445674%. SpatialBN
+              0 MB.          0%. Add
+              0 MB.          0%. Div
+              0 MB.          0%. Mul
+              0 MB.          0%. Relu
+        21.8994 MB in Total
+
+```
+### Optimized
+
+```
+Main run finished. Milliseconds per iter: 22.0981. Iters per second: 45.2527
+Time per operator type:
+         17.146 ms.    78.8965%. Conv
+        1.38453 ms.    6.37084%. Add
+        1.30991 ms.    6.02749%. Div
+       0.685417 ms.    3.15391%. Mul
+       0.532589 ms.    2.45068%. Relu
+       0.418263 ms.    1.92461%. Clip
+        0.15128 ms.   0.696106%. FC
+       0.102065 ms.   0.469648%. AveragePool
+      0.0022143 ms.   0.010189%. Squeeze
+        21.7323 ms in Total
+FLOP per operator type:
+       0.430616 GFLOP.    98.1927%. Conv
+       0.002561 GFLOP.   0.583981%. FC
+     0.00210961 GFLOP.   0.481051%. Mul
+     0.00173891 GFLOP.   0.396522%. Add
+     0.00151626 GFLOP.    0.34575%. Div
+              0 GFLOP.          0%. Relu
+       0.438542 GFLOP in Total
+Feature Memory Read per operator type:
+        34.7842 MB.     44.833%. Conv
+        14.5035 MB.    18.6934%. Mul
+        9.25778 MB.    11.9323%. Relu
+        7.84641 MB.    10.1132%. Add
+        6.06516 MB.    7.81733%. Div
+        5.12912 MB.    6.61087%. FC
+        77.5861 MB in Total
+Feature Memory Written per operator type:
+        17.6246 MB.    36.4556%. Conv
+        9.25778 MB.    19.1492%. Relu
+        8.43843 MB.    17.4544%. Mul
+        6.95565 MB.    14.3874%. Add
+        6.06502 MB.    12.5452%. Div
+          0.004 MB. 0.00827378%. FC
+        48.3455 MB in Total
+Parameter Memory per operator type:
+        16.6778 MB.    76.4973%. Conv
+          5.124 MB.    23.5027%. FC
+              0 MB.          0%. Add
+              0 MB.          0%. Div
+              0 MB.          0%. Mul
+              0 MB.          0%. Relu
+        21.8018 MB in Total
+
+```
+
+## MnasNet-A1
+
+### Unoptimized
+```
+Main run finished. Milliseconds per iter: 30.0892. Iters per second: 33.2345
+Time per operator type:
+        24.4656 ms.    79.0905%. Conv
+        4.14958 ms.    13.4144%. SpatialBN
+        1.60598 ms.    5.19169%. Relu
+       0.295219 ms.    0.95436%. Mul
+       0.187609 ms.   0.606486%. FC
+       0.120556 ms.   0.389724%. AveragePool
+        0.09036 ms.   0.292109%. Add
+       0.015727 ms.   0.050841%. Sigmoid
+     0.00306205 ms. 0.00989875%. Squeeze
+        30.9337 ms in Total
+FLOP per operator type:
+       0.620598 GFLOP.    95.6434%. Conv
+      0.0248873 GFLOP.     3.8355%. SpatialBN
+       0.002561 GFLOP.   0.394688%. FC
+    0.000597408 GFLOP.  0.0920695%. Mul
+    0.000222656 GFLOP.  0.0343146%. Add
+              0 GFLOP.          0%. Relu
+       0.648867 GFLOP in Total
+Feature Memory Read per operator type:
+        35.5457 MB.    38.4109%. Conv
+        25.1552 MB.    27.1829%. SpatialBN
+        22.5235 MB.     24.339%. Relu
+        5.12912 MB.    5.54256%. FC
+        2.40586 MB.    2.59978%. Mul
+        1.78125 MB.    1.92483%. Add
+        92.5406 MB in Total
+Feature Memory Written per operator type:
+        24.9042 MB.    32.9424%. Conv
+        24.8873 MB.      32.92%. SpatialBN
+        22.5235 MB.    29.7932%. Relu
+        2.38963 MB.    3.16092%. Mul
+       0.890624 MB.    1.17809%. Add
+          0.004 MB. 0.00529106%. FC
+        75.5993 MB in Total
+Parameter Memory per operator type:
+        10.2732 MB.    66.1459%. Conv
+          5.124 MB.    32.9917%. FC
+       0.133952 MB.    0.86247%. SpatialBN
+              0 MB.          0%. Add
+              0 MB.          0%. Mul
+              0 MB.          0%. Relu
+        15.5312 MB in Total
+```
+
+### Optimized
+```
+Main run finished. Milliseconds per iter: 24.2367. Iters per second: 41.2597
+Time per operator type:
+        22.0547 ms.    91.1375%. Conv
+        1.49096 ms.    6.16116%. Relu
+       0.253417 ms.     1.0472%. Mul
+        0.18506 ms.    0.76473%. FC
+       0.112942 ms.   0.466717%. AveragePool
+       0.086769 ms.   0.358559%. Add
+      0.0127889 ms.  0.0528479%. Sigmoid
+      0.0027346 ms.  0.0113003%. Squeeze
+        24.1994 ms in Total
+FLOP per operator type:
+       0.620598 GFLOP.    99.4581%. Conv
+       0.002561 GFLOP.    0.41043%. FC
+    0.000597408 GFLOP.  0.0957417%. Mul
+    0.000222656 GFLOP.  0.0356832%. Add
+              0 GFLOP.          0%. Relu
+       0.623979 GFLOP in Total
+Feature Memory Read per operator type:
+        35.6127 MB.    52.7968%. Conv
+        22.5235 MB.    33.3917%. Relu
+        5.12912 MB.    7.60406%. FC
+        2.40586 MB.    3.56675%. Mul
+        1.78125 MB.    2.64075%. Add
+        67.4524 MB in Total
+Feature Memory Written per operator type:
+        24.9042 MB.    49.1092%. Conv
+        22.5235 MB.    44.4145%. Relu
+        2.38963 MB.    4.71216%. Mul
+       0.890624 MB.    1.75624%. Add
+          0.004 MB. 0.00788768%. FC
+         50.712 MB in Total
+Parameter Memory per operator type:
+        10.2732 MB.    66.7213%. Conv
+          5.124 MB.    33.2787%. FC
+              0 MB.          0%. Add
+              0 MB.          0%. Mul
+              0 MB.          0%. Relu
+        15.3972 MB in Total
+```
+## MnasNet-B1
+
+### Unoptimized
+```
+Main run finished. Milliseconds per iter: 28.3109. Iters per second: 35.322
+Time per operator type:
+        29.1121 ms.    83.3081%. Conv
+        4.14959 ms.    11.8746%. SpatialBN
+        1.35823 ms.    3.88675%. Relu
+       0.186188 ms.   0.532802%. FC
+       0.116244 ms.   0.332647%. Add
+       0.018641 ms.  0.0533437%. AveragePool
+      0.0040904 ms.  0.0117052%. Squeeze
+        34.9451 ms in Total
+FLOP per operator type:
+       0.626272 GFLOP.    96.2088%. Conv
+      0.0218266 GFLOP.    3.35303%. SpatialBN
+       0.002561 GFLOP.   0.393424%. FC
+    0.000291648 GFLOP.  0.0448034%. Add
+              0 GFLOP.          0%. Relu
+       0.650951 GFLOP in Total
+Feature Memory Read per operator type:
+        34.4354 MB.    41.3788%. Conv
+        22.1299 MB.    26.5921%. SpatialBN
+        19.1923 MB.    23.0622%. Relu
+        5.12912 MB.    6.16333%. FC
+        2.33318 MB.    2.80364%. Add
+        83.2199 MB in Total
+Feature Memory Written per operator type:
+        21.8266 MB.    34.0955%. Conv
+        21.8266 MB.    34.0955%. SpatialBN
+        19.1923 MB.    29.9805%. Relu
+        1.16659 MB.    1.82234%. Add
+          0.004 MB. 0.00624844%. FC
+         64.016 MB in Total
+Parameter Memory per operator type:
+        12.2576 MB.    69.9104%. Conv
+          5.124 MB.    29.2245%. FC
+        0.15168 MB.   0.865099%. SpatialBN
+              0 MB.          0%. Add
+              0 MB.          0%. Relu
+        17.5332 MB in Total
+```
+
+### Optimized
+```
+Main run finished. Milliseconds per iter: 26.6364. Iters per second: 37.5426
+Time per operator type:
+        24.9888 ms.    94.0962%. Conv
+        1.26147 ms.    4.75011%. Relu
+       0.176234 ms.   0.663619%. FC
+       0.113309 ms.   0.426672%. Add
+      0.0138708 ms.  0.0522311%. AveragePool
+     0.00295685 ms.  0.0111341%. Squeeze
+        26.5566 ms in Total
+FLOP per operator type:
+       0.626272 GFLOP.    99.5466%. Conv
+       0.002561 GFLOP.   0.407074%. FC
+    0.000291648 GFLOP.  0.0463578%. Add
+              0 GFLOP.          0%. Relu
+       0.629124 GFLOP in Total
+Feature Memory Read per operator type:
+        34.5112 MB.    56.4224%. Conv
+        19.1923 MB.    31.3775%. Relu
+        5.12912 MB.     8.3856%. FC
+        2.33318 MB.    3.81452%. Add
+        61.1658 MB in Total
+Feature Memory Written per operator type:
+        21.8266 MB.    51.7346%. Conv
+        19.1923 MB.    45.4908%. Relu
+        1.16659 MB.    2.76513%. Add
+          0.004 MB. 0.00948104%. FC
+        42.1895 MB in Total
+Parameter Memory per operator type:
+        12.2576 MB.    70.5205%. Conv
+          5.124 MB.    29.4795%. FC
+              0 MB.          0%. Add
+              0 MB.          0%. Relu
+        17.3816 MB in Total
+```

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "{}"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2020 Ross Wightman
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/README.md

@@ -0,0 +1,323 @@
+# (Generic) EfficientNets for PyTorch
+
+A 'generic' implementation of EfficientNet, MixNet, MobileNetV3, etc. that covers most of the compute/parameter efficient architectures derived from the MobileNet V1/V2 block sequence, including those found via automated neural architecture search. 
+
+All models are implemented by GenEfficientNet or MobileNetV3 classes, with string based architecture definitions to configure the block layouts (idea from [here](https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mnasnet_models.py))
+
+## What's New
+
+### Aug 19, 2020
+* Add updated PyTorch trained EfficientNet-B3 weights trained by myself with `timm` (82.1 top-1)
+* Add PyTorch trained EfficientNet-Lite0 contributed by [@hal-314](https://github.com/hal-314) (75.5 top-1)
+* Update ONNX and Caffe2 export / utility scripts to work with latest PyTorch / ONNX
+* ONNX runtime based validation script added
+* activations (mostly) brought in sync with `timm` equivalents
+
+
+### April 5, 2020
+* Add some newly trained MobileNet-V2 models trained with latest h-params, rand augment. They compare quite favourably to EfficientNet-Lite
+  * 3.5M param MobileNet-V2 100 @ 73%
+  * 4.5M param MobileNet-V2 110d @ 75%
+  * 6.1M param MobileNet-V2 140 @ 76.5%
+  * 5.8M param MobileNet-V2 120d @ 77.3%
+  
+### March 23, 2020
+ * Add EfficientNet-Lite models w/ weights ported from [Tensorflow TPU](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/lite)
+ * Add PyTorch trained MobileNet-V3 Large weights with 75.77% top-1
+ * IMPORTANT CHANGE (if training from scratch) - weight init changed to better match Tensorflow impl, set `fix_group_fanout=False` in `initialize_weight_goog` for old behavior
+
+### Feb 12, 2020
+ * Add EfficientNet-L2 and B0-B7 NoisyStudent weights ported from [Tensorflow TPU](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet)
+ * Port new EfficientNet-B8 (RandAugment) weights from TF TPU, these are different than the B8 AdvProp, different input normalization.
+ * Add RandAugment PyTorch trained EfficientNet-ES (EdgeTPU-Small) weights with 78.1 top-1. Trained by [Andrew Lavin](https://github.com/andravin)
+
+### Jan 22, 2020
+ * Update weights for EfficientNet B0, B2, B3 and MixNet-XL with latest RandAugment trained weights. Trained with (https://github.com/rwightman/pytorch-image-models)
+ * Fix torchscript compatibility for PyTorch 1.4, add torchscript support for MixedConv2d using ModuleDict
+ * Test models, torchscript, onnx export with PyTorch 1.4 -- no issues
+
+### Nov 22, 2019
+ * New top-1 high! Ported official TF EfficientNet AdvProp (https://arxiv.org/abs/1911.09665) weights and B8 model spec. Created a new set of `ap` models since they use a different
+ preprocessing (Inception mean/std) from the original EfficientNet base/AA/RA weights.
+ 
+### Nov 15, 2019
+ * Ported official TF MobileNet-V3 float32 large/small/minimalistic weights
+ * Modifications to MobileNet-V3 model and components to support some additional config needed for differences between TF MobileNet-V3 and mine
+
+### Oct 30, 2019
+ * Many of the models will now work with torch.jit.script, MixNet being the biggest exception
+ * Improved interface for enabling torchscript or ONNX export compatible modes (via config)
+ * Add JIT optimized mem-efficient Swish/Mish autograd.fn in addition to memory-efficient autgrad.fn
+ * Activation factory to select best version of activation by name or override one globally
+ * Add pretrained checkpoint load helper that handles input conv and classifier changes
+ 
+### Oct 27, 2019
+ * Add CondConv EfficientNet variants ported from https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/condconv
+ * Add RandAug weights for TF EfficientNet B5 and B7 from https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet
+ * Bring over MixNet-XL model and depth scaling algo from my pytorch-image-models code base
+ * Switch activations and global pooling to modules
+ * Add memory-efficient Swish/Mish impl
+ * Add as_sequential() method to all models and allow as an argument in entrypoint fns
+ * Move MobileNetV3 into own file since it has a different head
+ * Remove ChamNet, MobileNet V2/V1 since they will likely never be used here
+
+## Models
+
+Implemented models include:
+  * EfficientNet NoisyStudent (B0-B7, L2) (https://arxiv.org/abs/1911.04252)
+  * EfficientNet AdvProp (B0-B8) (https://arxiv.org/abs/1911.09665)
+  * EfficientNet (B0-B8) (https://arxiv.org/abs/1905.11946)
+  * EfficientNet-EdgeTPU (S, M, L) (https://ai.googleblog.com/2019/08/efficientnet-edgetpu-creating.html)
+  * EfficientNet-CondConv (https://arxiv.org/abs/1904.04971)
+  * EfficientNet-Lite (https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/lite)
+  * MixNet (https://arxiv.org/abs/1907.09595)
+  * MNASNet B1, A1 (Squeeze-Excite), and Small (https://arxiv.org/abs/1807.11626)
+  * MobileNet-V3 (https://arxiv.org/abs/1905.02244)
+  * FBNet-C (https://arxiv.org/abs/1812.03443)
+  * Single-Path NAS (https://arxiv.org/abs/1904.02877)
+    
+I originally implemented and trained some these models with code [here](https://github.com/rwightman/pytorch-image-models), this repository contains just the GenEfficientNet models, validation, and associated ONNX/Caffe2 export code. 
+
+## Pretrained
+
+I've managed to train several of the models to accuracies close to or above the originating papers and official impl. My training code is here: https://github.com/rwightman/pytorch-image-models
+
+
+|Model | Prec@1 (Err) | Prec@5 (Err) | Param#(M) | MAdds(M) | Image Scaling | Resolution | Crop |
+|---|---|---|---|---|---|---|---|
+| efficientnet_b3 | 82.240 (17.760) | 96.116 (3.884) | 12.23 | TBD | bicubic | 320 | 1.0 |
+| efficientnet_b3 | 82.076 (17.924) | 96.020 (3.980) | 12.23 | TBD | bicubic | 300 | 0.904 |
+| mixnet_xl | 81.074 (18.926) | 95.282 (4.718) | 11.90 | TBD | bicubic | 256 | 1.0 |
+| efficientnet_b2 | 80.612 (19.388) | 95.318 (4.682) | 9.1 | TBD | bicubic | 288 | 1.0 |
+| mixnet_xl | 80.476 (19.524) | 94.936 (5.064) | 11.90 | TBD | bicubic | 224 | 0.875 |
+| efficientnet_b2 | 80.288 (19.712) | 95.166 (4.834) | 9.1 | 1003 | bicubic | 260 | 0.890 |
+| mixnet_l | 78.976 (21.024 | 94.184 (5.816) | 7.33 | TBD | bicubic | 224 | 0.875 |
+| efficientnet_b1 | 78.692 (21.308) | 94.086 (5.914) | 7.8 | 694 | bicubic | 240 | 0.882 |
+| efficientnet_es | 78.066 (21.934) | 93.926 (6.074) | 5.44 | TBD | bicubic | 224 | 0.875 |
+| efficientnet_b0 | 77.698 (22.302) | 93.532 (6.468) | 5.3 | 390 | bicubic | 224 | 0.875 |
+| mobilenetv2_120d | 77.294 (22.706 | 93.502 (6.498) | 5.8 | TBD | bicubic | 224 | 0.875 |
+| mixnet_m | 77.256 (22.744) | 93.418 (6.582) | 5.01 | 353 | bicubic | 224 | 0.875 |
+| mobilenetv2_140 | 76.524 (23.476) | 92.990 (7.010) | 6.1 | TBD | bicubic | 224 | 0.875 |
+| mixnet_s | 75.988 (24.012) | 92.794 (7.206) | 4.13 | TBD | bicubic | 224 | 0.875 |
+| mobilenetv3_large_100 | 75.766 (24.234) | 92.542 (7.458) | 5.5 | TBD | bicubic | 224 | 0.875 |
+| mobilenetv3_rw | 75.634 (24.366) | 92.708 (7.292) | 5.5 | 219 | bicubic | 224 | 0.875 |
+| efficientnet_lite0 | 75.472 (24.528) | 92.520 (7.480) | 4.65 | TBD | bicubic | 224 | 0.875 |
+| mnasnet_a1 | 75.448 (24.552) | 92.604 (7.396) | 3.9 | 312 | bicubic | 224 | 0.875 |
+| fbnetc_100 | 75.124 (24.876) | 92.386 (7.614) | 5.6 | 385 | bilinear | 224 | 0.875 |
+| mobilenetv2_110d | 75.052 (24.948) | 92.180 (7.820) | 4.5 | TBD | bicubic | 224 | 0.875 |
+| mnasnet_b1 | 74.658 (25.342) | 92.114 (7.886) | 4.4 | 315 | bicubic | 224 | 0.875 |
+| spnasnet_100 | 74.084 (25.916)  | 91.818 (8.182) | 4.4 | TBD | bilinear | 224 | 0.875 |
+| mobilenetv2_100 | 72.978 (27.022) | 91.016 (8.984) | 3.5 | TBD | bicubic | 224 | 0.875 |
+
+
+More pretrained models to come...
+
+
+## Ported Weights
+
+The weights ported from Tensorflow checkpoints for the EfficientNet models do pretty much match accuracy in Tensorflow once a SAME convolution padding equivalent is added, and the same crop factors, image scaling, etc (see table) are used via cmd line args.
+
+**IMPORTANT:** 
+* Tensorflow ported weights for EfficientNet AdvProp (AP), EfficientNet EdgeTPU, EfficientNet-CondConv, EfficientNet-Lite, and MobileNet-V3 models use Inception style (0.5, 0.5, 0.5) for mean and std.
+* Enabling the Tensorflow preprocessing pipeline with `--tf-preprocessing` at validation time will improve scores by 0.1-0.5%, very close to original TF impl.
+
+To run validation for tf_efficientnet_b5:
+`python validate.py /path/to/imagenet/validation/ --model tf_efficientnet_b5 -b 64 --img-size 456 --crop-pct 0.934 --interpolation bicubic`
+
+To run validation w/ TF preprocessing for tf_efficientnet_b5:
+`python validate.py /path/to/imagenet/validation/ --model tf_efficientnet_b5 -b 64 --img-size 456 --tf-preprocessing`
+
+To run validation for a model with Inception preprocessing, ie EfficientNet-B8 AdvProp:
+`python validate.py /path/to/imagenet/validation/ --model tf_efficientnet_b8_ap -b 48 --num-gpu 2 --img-size 672 --crop-pct 0.954 --mean 0.5 --std 0.5`
+
+|Model | Prec@1 (Err) | Prec@5 (Err) | Param # | Image Scaling  | Image Size | Crop | 
+|---|---|---|---|---|---|---|
+| tf_efficientnet_l2_ns *tfp | 88.352 (11.648) | 98.652 (1.348) | 480 | bicubic | 800 | N/A |
+| tf_efficientnet_l2_ns      | TBD | TBD | 480 | bicubic | 800 | 0.961 |
+| tf_efficientnet_l2_ns_475      | 88.234 (11.766) | 98.546 (1.454) | 480 | bicubic | 475 | 0.936 |
+| tf_efficientnet_l2_ns_475 *tfp | 88.172 (11.828) | 98.566 (1.434) | 480 | bicubic | 475 | N/A |
+| tf_efficientnet_b7_ns *tfp | 86.844 (13.156) | 98.084 (1.916) | 66.35 | bicubic | 600 | N/A |
+| tf_efficientnet_b7_ns      | 86.840 (13.160) | 98.094 (1.906) | 66.35 | bicubic | 600 | N/A |
+| tf_efficientnet_b6_ns      | 86.452 (13.548) | 97.882 (2.118) | 43.04 | bicubic | 528 | N/A |
+| tf_efficientnet_b6_ns *tfp | 86.444 (13.556) | 97.880 (2.120) | 43.04 | bicubic | 528 | N/A |
+| tf_efficientnet_b5_ns *tfp | 86.064 (13.936) | 97.746 (2.254) | 30.39 | bicubic | 456 | N/A |
+| tf_efficientnet_b5_ns      | 86.088 (13.912) | 97.752 (2.248) | 30.39 | bicubic | 456 | N/A |
+| tf_efficientnet_b8_ap *tfp | 85.436 (14.564) | 97.272 (2.728) | 87.4 | bicubic | 672 | N/A |
+| tf_efficientnet_b8 *tfp    | 85.384 (14.616) | 97.394 (2.606) | 87.4 | bicubic | 672 | N/A |
+| tf_efficientnet_b8         | 85.370 (14.630) | 97.390 (2.610) | 87.4 | bicubic | 672 | 0.954 |
+| tf_efficientnet_b8_ap      | 85.368 (14.632) | 97.294 (2.706) | 87.4 | bicubic | 672 | 0.954 |
+| tf_efficientnet_b4_ns *tfp | 85.298 (14.702) | 97.504 (2.496) | 19.34 | bicubic | 380 | N/A |
+| tf_efficientnet_b4_ns      | 85.162 (14.838) | 97.470 (2.530) | 19.34 | bicubic | 380 | 0.922 |
+| tf_efficientnet_b7_ap *tfp | 85.154 (14.846) | 97.244 (2.756) | 66.35 | bicubic | 600 | N/A |
+| tf_efficientnet_b7_ap      | 85.118 (14.882) | 97.252 (2.748) | 66.35 | bicubic | 600 | 0.949 |
+| tf_efficientnet_b7 *tfp | 84.940 (15.060) | 97.214 (2.786) | 66.35 | bicubic | 600 | N/A |
+| tf_efficientnet_b7      | 84.932 (15.068) | 97.208 (2.792) | 66.35 | bicubic | 600 | 0.949 |
+| tf_efficientnet_b6_ap      | 84.786 (15.214) | 97.138 (2.862) | 43.04 | bicubic | 528 | 0.942 |
+| tf_efficientnet_b6_ap *tfp | 84.760 (15.240) | 97.124 (2.876) | 43.04 | bicubic | 528 | N/A |
+| tf_efficientnet_b5_ap *tfp | 84.276 (15.724) | 96.932 (3.068) | 30.39 | bicubic | 456 | N/A |
+| tf_efficientnet_b5_ap      | 84.254 (15.746) | 96.976 (3.024) | 30.39 | bicubic | 456 | 0.934 |
+| tf_efficientnet_b6 *tfp  | 84.140 (15.860) | 96.852 (3.148) | 43.04 | bicubic | 528 | N/A |
+| tf_efficientnet_b6       | 84.110 (15.890) | 96.886 (3.114) | 43.04 | bicubic | 528 | 0.942 |
+| tf_efficientnet_b3_ns *tfp | 84.054 (15.946) | 96.918 (3.082) | 12.23 | bicubic | 300 | N/A |
+| tf_efficientnet_b3_ns      | 84.048 (15.952) | 96.910 (3.090) | 12.23 | bicubic | 300 | .904 |
+| tf_efficientnet_b5 *tfp  | 83.822 (16.178) | 96.756 (3.244) | 30.39 | bicubic | 456 | N/A |
+| tf_efficientnet_b5       | 83.812 (16.188) | 96.748 (3.252) | 30.39 | bicubic | 456 | 0.934 |
+| tf_efficientnet_b4_ap *tfp | 83.278 (16.722) | 96.376 (3.624) | 19.34 | bicubic | 380 | N/A |
+| tf_efficientnet_b4_ap      | 83.248 (16.752) | 96.388 (3.612) | 19.34 | bicubic | 380 | 0.922 |
+| tf_efficientnet_b4       | 83.022 (16.978) | 96.300 (3.700) | 19.34 | bicubic | 380 | 0.922 |
+| tf_efficientnet_b4 *tfp  | 82.948 (17.052) | 96.308 (3.692) | 19.34 | bicubic | 380 | N/A |
+| tf_efficientnet_b2_ns *tfp | 82.436 (17.564) | 96.268 (3.732) | 9.11 | bicubic | 260 | N/A |
+| tf_efficientnet_b2_ns      | 82.380 (17.620) | 96.248 (3.752) | 9.11 | bicubic | 260 | 0.89 |
+| tf_efficientnet_b3_ap *tfp | 81.882 (18.118) | 95.662 (4.338) | 12.23 | bicubic | 300 | N/A |
+| tf_efficientnet_b3_ap      | 81.828 (18.172) | 95.624 (4.376) | 12.23 | bicubic | 300 | 0.904 |
+| tf_efficientnet_b3       | 81.636 (18.364) | 95.718 (4.282) | 12.23 | bicubic | 300 | 0.904 |
+| tf_efficientnet_b3 *tfp  | 81.576 (18.424) | 95.662 (4.338) | 12.23 | bicubic | 300 | N/A |
+| tf_efficientnet_lite4      | 81.528 (18.472) | 95.668 (4.332) | 13.00  | bilinear | 380 | 0.92 |
+| tf_efficientnet_b1_ns *tfp | 81.514 (18.486) | 95.776 (4.224) | 7.79 | bicubic | 240 | N/A |
+| tf_efficientnet_lite4 *tfp | 81.502 (18.498) | 95.676 (4.324) | 13.00  | bilinear | 380 | N/A |
+| tf_efficientnet_b1_ns      | 81.388 (18.612) | 95.738 (4.262) | 7.79 | bicubic | 240 | 0.88 |
+| tf_efficientnet_el       | 80.534 (19.466) | 95.190 (4.810) | 10.59 | bicubic | 300 | 0.904 |
+| tf_efficientnet_el *tfp  | 80.476 (19.524) | 95.200 (4.800) | 10.59 | bicubic | 300 | N/A |
+| tf_efficientnet_b2_ap *tfp | 80.420 (19.580) | 95.040 (4.960) | 9.11 | bicubic | 260 | N/A |
+| tf_efficientnet_b2_ap    | 80.306 (19.694) | 95.028 (4.972) | 9.11 | bicubic | 260 | 0.890 |
+| tf_efficientnet_b2 *tfp  | 80.188 (19.812) | 94.974 (5.026) | 9.11 | bicubic | 260 | N/A |
+| tf_efficientnet_b2       | 80.086 (19.914) | 94.908 (5.092) | 9.11 | bicubic | 260 | 0.890 |
+| tf_efficientnet_lite3       | 79.812 (20.188) | 94.914 (5.086) | 8.20  | bilinear | 300 | 0.904 |
+| tf_efficientnet_lite3 *tfp  | 79.734 (20.266) | 94.838 (5.162) | 8.20  | bilinear | 300 | N/A |
+| tf_efficientnet_b1_ap *tfp | 79.532 (20.468) | 94.378 (5.622) | 7.79 | bicubic | 240 | N/A |
+| tf_efficientnet_cc_b1_8e *tfp | 79.464 (20.536)| 94.492 (5.508) | 39.7 | bicubic | 240 | 0.88 |
+| tf_efficientnet_cc_b1_8e | 79.298 (20.702) | 94.364 (5.636) | 39.7 | bicubic | 240 | 0.88 |
+| tf_efficientnet_b1_ap    | 79.278 (20.722) | 94.308 (5.692) | 7.79 | bicubic | 240 | 0.88 |
+| tf_efficientnet_b1 *tfp  | 79.172 (20.828) | 94.450 (5.550) | 7.79 | bicubic | 240 | N/A |
+| tf_efficientnet_em *tfp  | 78.958 (21.042) | 94.458 (5.542) | 6.90 | bicubic | 240 | N/A |
+| tf_efficientnet_b0_ns *tfp | 78.806 (21.194) | 94.496 (5.504) | 5.29 | bicubic | 224 | N/A |
+| tf_mixnet_l *tfp         | 78.846 (21.154) | 94.212 (5.788) | 7.33 | bilinear | 224 | N/A |
+| tf_efficientnet_b1       | 78.826 (21.174) | 94.198 (5.802) | 7.79 | bicubic | 240 | 0.88 |
+| tf_mixnet_l              | 78.770 (21.230) | 94.004 (5.996) | 7.33 | bicubic | 224 | 0.875 |
+| tf_efficientnet_em       | 78.742 (21.258) | 94.332 (5.668) | 6.90 | bicubic | 240 | 0.875 |
+| tf_efficientnet_b0_ns    | 78.658 (21.342) | 94.376 (5.624) | 5.29 | bicubic | 224 | 0.875 |
+| tf_efficientnet_cc_b0_8e *tfp | 78.314 (21.686) | 93.790 (6.210) | 24.0 | bicubic | 224 | 0.875 |
+| tf_efficientnet_cc_b0_8e | 77.908 (22.092) | 93.656 (6.344) | 24.0 | bicubic | 224 | 0.875 |
+| tf_efficientnet_cc_b0_4e *tfp | 77.746 (22.254) | 93.552 (6.448) | 13.3 | bicubic | 224 | 0.875 |
+| tf_efficientnet_cc_b0_4e | 77.304 (22.696) | 93.332 (6.668) | 13.3 | bicubic | 224 | 0.875 |
+| tf_efficientnet_es *tfp  | 77.616 (22.384) | 93.750 (6.250) | 5.44 | bicubic | 224 | N/A |
+| tf_efficientnet_lite2 *tfp  | 77.544 (22.456) | 93.800 (6.200) | 6.09  | bilinear | 260 | N/A |
+| tf_efficientnet_lite2       | 77.460 (22.540) | 93.746 (6.254) | 6.09  | bicubic | 260 | 0.89 |
+| tf_efficientnet_b0_ap *tfp | 77.514 (22.486) | 93.576 (6.424) | 5.29  | bicubic | 224 | N/A |
+| tf_efficientnet_es       | 77.264 (22.736) | 93.600 (6.400) | 5.44 | bicubic | 224 | N/A |
+| tf_efficientnet_b0 *tfp  | 77.258 (22.742) | 93.478 (6.522) | 5.29  | bicubic | 224 | N/A |
+| tf_efficientnet_b0_ap    | 77.084 (22.916) | 93.254 (6.746) | 5.29  | bicubic | 224 | 0.875 |
+| tf_mixnet_m *tfp         | 77.072 (22.928) | 93.368 (6.632) | 5.01 | bilinear | 224 | N/A |
+| tf_mixnet_m              | 76.950 (23.050) | 93.156 (6.844) | 5.01 | bicubic | 224 | 0.875 |
+| tf_efficientnet_b0       | 76.848 (23.152) | 93.228 (6.772) | 5.29  | bicubic | 224 | 0.875 |
+| tf_efficientnet_lite1 *tfp  | 76.764 (23.236) | 93.326 (6.674) | 5.42  | bilinear | 240 | N/A |
+| tf_efficientnet_lite1       | 76.638 (23.362) | 93.232 (6.768) | 5.42  | bicubic | 240 | 0.882 |
+| tf_mixnet_s *tfp         | 75.800 (24.200) | 92.788 (7.212) | 4.13 | bilinear | 224 | N/A |
+| tf_mobilenetv3_large_100 *tfp | 75.768 (24.232) | 92.710 (7.290) | 5.48 | bilinear | 224 | N/A |
+| tf_mixnet_s              | 75.648 (24.352) | 92.636 (7.364) | 4.13 | bicubic | 224 | 0.875 |
+| tf_mobilenetv3_large_100 | 75.516 (24.484) | 92.600 (7.400) | 5.48 | bilinear | 224 | 0.875 |
+| tf_efficientnet_lite0 *tfp  | 75.074 (24.926) | 92.314 (7.686) | 4.65  | bilinear | 224 | N/A |
+| tf_efficientnet_lite0       | 74.842 (25.158) | 92.170 (7.830) | 4.65  | bicubic | 224 | 0.875 |
+| tf_mobilenetv3_large_075 *tfp | 73.730 (26.270) | 91.616 (8.384) | 3.99 | bilinear | 224 |N/A |
+| tf_mobilenetv3_large_075 | 73.442 (26.558) | 91.352 (8.648) | 3.99 | bilinear | 224 | 0.875 |
+| tf_mobilenetv3_large_minimal_100 *tfp | 72.678 (27.322) | 90.860 (9.140) | 3.92 | bilinear | 224 | N/A |
+| tf_mobilenetv3_large_minimal_100 | 72.244 (27.756) | 90.636 (9.364) | 3.92 | bilinear | 224 | 0.875 |
+| tf_mobilenetv3_small_100 *tfp | 67.918 (32.082) | 87.958 (12.042 | 2.54 | bilinear | 224 | N/A |
+| tf_mobilenetv3_small_100 | 67.918 (32.082) | 87.662 (12.338) | 2.54 | bilinear | 224 | 0.875 |
+| tf_mobilenetv3_small_075 *tfp | 66.142 (33.858) | 86.498 (13.502) | 2.04 | bilinear | 224 | N/A |
+| tf_mobilenetv3_small_075 | 65.718 (34.282) | 86.136 (13.864) | 2.04 | bilinear | 224 | 0.875 |
+| tf_mobilenetv3_small_minimal_100 *tfp | 63.378 (36.622) | 84.802 (15.198) | 2.04 | bilinear | 224 | N/A |
+| tf_mobilenetv3_small_minimal_100 | 62.898 (37.102) | 84.230 (15.770) | 2.04 | bilinear | 224 | 0.875 |
+
+
+*tfp models validated with `tf-preprocessing` pipeline
+
+Google tf and tflite weights ported from official Tensorflow repositories
+* https://github.com/tensorflow/tpu/tree/master/models/official/mnasnet
+* https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet
+* https://github.com/tensorflow/models/tree/master/research/slim/nets/mobilenet
+
+## Usage
+
+### Environment
+
+All development and testing has been done in Conda Python 3 environments on Linux x86-64 systems, specifically Python 3.6.x, 3.7.x, 3.8.x.
+
+Users have reported that a Python 3 Anaconda install in Windows works. I have not verified this myself.
+
+PyTorch versions 1.4, 1.5, 1.6 have been tested with this code.
+
+I've tried to keep the dependencies minimal, the setup is as per the PyTorch default install instructions for Conda:
+```
+conda create -n torch-env
+conda activate torch-env
+conda install -c pytorch pytorch torchvision cudatoolkit=10.2
+```
+
+### PyTorch Hub
+
+Models can be accessed via the PyTorch Hub API
+
+```
+>>> torch.hub.list('rwightman/gen-efficientnet-pytorch')
+['efficientnet_b0', ...]
+>>> model = torch.hub.load('rwightman/gen-efficientnet-pytorch', 'efficientnet_b0', pretrained=True)
+>>> model.eval()
+>>> output = model(torch.randn(1,3,224,224))
+```
+
+### Pip
+This package can be installed via pip.
+
+Install (after conda env/install):
+```
+pip install geffnet
+```
+
+Eval use:
+```
+>>> import geffnet
+>>> m = geffnet.create_model('mobilenetv3_large_100', pretrained=True)
+>>> m.eval()
+```
+
+Train use:
+```
+>>> import geffnet
+>>> # models can also be created by using the entrypoint directly
+>>> m = geffnet.efficientnet_b2(pretrained=True, drop_rate=0.25, drop_connect_rate=0.2)
+>>> m.train()
+```
+
+Create in a nn.Sequential container, for fast.ai, etc:
+```
+>>> import geffnet
+>>> m = geffnet.mixnet_l(pretrained=True, drop_rate=0.25, drop_connect_rate=0.2, as_sequential=True)
+```
+
+### Exporting
+
+Scripts are included to
+* export models to ONNX (`onnx_export.py`)
+* optimized ONNX graph (`onnx_optimize.py` or `onnx_validate.py` w/ `--onnx-output-opt` arg)
+* validate with ONNX runtime  (`onnx_validate.py`)
+* convert ONNX model to Caffe2 (`onnx_to_caffe.py`)
+* validate in Caffe2 (`caffe2_validate.py`)
+* benchmark in Caffe2 w/ FLOPs, parameters output (`caffe2_benchmark.py`)
+
+As an example, to export the MobileNet-V3 pretrained model and then run an Imagenet validation:
+```
+python onnx_export.py --model mobilenetv3_large_100 ./mobilenetv3_100.onnx
+python onnx_validate.py /imagenet/validation/ --onnx-input ./mobilenetv3_100.onnx 
+```
+
+These scripts were tested to be working as of PyTorch 1.6 and ONNX 1.7 w/ ONNX runtime 1.4. Caffe2 compatible
+export now requires additional args mentioned in the export script (not needed in earlier versions).
+
+#### Export Notes
+1. The TF ported weights with the 'SAME' conv padding activated cannot be exported to ONNX unless `_EXPORTABLE` flag in `config.py` is set to True. Use `config.set_exportable(True)` as in the `onnx_export.py` script.
+2. TF ported models with 'SAME' padding will have the padding fixed at export time to the resolution used for export. Even though dynamic padding is supported in opset >= 11, I can't get it working.
+3. ONNX optimize facility doesn't work reliably in PyTorch 1.6 / ONNX 1.7. Fortunately, the onnxruntime based inference is working very well now and includes on the fly optimization.
+3. ONNX / Caffe2 export/import frequently breaks with different PyTorch and ONNX version releases. Please check their respective issue trackers before filing issues here.
+
+

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/caffe2_benchmark.py

@@ -0,0 +1,65 @@
+""" Caffe2 validation script
+
+This script runs Caffe2 benchmark on exported ONNX model.
+It is a useful tool for reporting model FLOPS.
+
+Copyright 2020 Ross Wightman
+"""
+import argparse
+from caffe2.python import core, workspace, model_helper
+from caffe2.proto import caffe2_pb2
+
+
+parser = argparse.ArgumentParser(description='Caffe2 Model Benchmark')
+parser.add_argument('--c2-prefix', default='', type=str, metavar='NAME',
+                    help='caffe2 model pb name prefix')
+parser.add_argument('--c2-init', default='', type=str, metavar='PATH',
+                    help='caffe2 model init .pb')
+parser.add_argument('--c2-predict', default='', type=str, metavar='PATH',
+                    help='caffe2 model predict .pb')
+parser.add_argument('-b', '--batch-size', default=1, type=int,
+                    metavar='N', help='mini-batch size (default: 1)')
+parser.add_argument('--img-size', default=224, type=int,
+                    metavar='N', help='Input image dimension, uses model default if empty')
+
+
+def main():
+    args = parser.parse_args()
+    args.gpu_id = 0
+    if args.c2_prefix:
+        args.c2_init = args.c2_prefix + '.init.pb'
+        args.c2_predict = args.c2_prefix + '.predict.pb'
+
+    model = model_helper.ModelHelper(name="le_net", init_params=False)
+
+    # Bring in the init net from init_net.pb
+    init_net_proto = caffe2_pb2.NetDef()
+    with open(args.c2_init, "rb") as f:
+        init_net_proto.ParseFromString(f.read())
+    model.param_init_net = core.Net(init_net_proto)
+
+    # bring in the predict net from predict_net.pb
+    predict_net_proto = caffe2_pb2.NetDef()
+    with open(args.c2_predict, "rb") as f:
+        predict_net_proto.ParseFromString(f.read())
+    model.net = core.Net(predict_net_proto)
+
+    # CUDA performance not impressive
+    #device_opts = core.DeviceOption(caffe2_pb2.PROTO_CUDA, args.gpu_id)
+    #model.net.RunAllOnGPU(gpu_id=args.gpu_id, use_cudnn=True)
+    #model.param_init_net.RunAllOnGPU(gpu_id=args.gpu_id, use_cudnn=True)
+
+    input_blob = model.net.external_inputs[0]
+    model.param_init_net.GaussianFill(
+        [],
+        input_blob.GetUnscopedName(),
+        shape=(args.batch_size, 3, args.img_size, args.img_size),
+        mean=0.0,
+        std=1.0)
+    workspace.RunNetOnce(model.param_init_net)
+    workspace.CreateNet(model.net, overwrite=True)
+    workspace.BenchmarkNet(model.net.Proto().name, 5, 20, True)
+
+
+if __name__ == '__main__':
+    main()

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/caffe2_validate.py

@@ -0,0 +1,138 @@
+""" Caffe2 validation script
+
+This script is created to verify exported ONNX models running in Caffe2
+It utilizes the same PyTorch dataloader/processing pipeline for a
+fair comparison against the originals.
+
+Copyright 2020 Ross Wightman
+"""
+import argparse
+import numpy as np
+from caffe2.python import core, workspace, model_helper
+from caffe2.proto import caffe2_pb2
+from data import create_loader, resolve_data_config, Dataset
+from utils import AverageMeter
+import time
+
+parser = argparse.ArgumentParser(description='Caffe2 ImageNet Validation')
+parser.add_argument('data', metavar='DIR',
+                    help='path to dataset')
+parser.add_argument('--c2-prefix', default='', type=str, metavar='NAME',
+                    help='caffe2 model pb name prefix')
+parser.add_argument('--c2-init', default='', type=str, metavar='PATH',
+                    help='caffe2 model init .pb')
+parser.add_argument('--c2-predict', default='', type=str, metavar='PATH',
+                    help='caffe2 model predict .pb')
+parser.add_argument('-j', '--workers', default=2, type=int, metavar='N',
+                    help='number of data loading workers (default: 2)')
+parser.add_argument('-b', '--batch-size', default=256, type=int,
+                    metavar='N', help='mini-batch size (default: 256)')
+parser.add_argument('--img-size', default=None, type=int,
+                    metavar='N', help='Input image dimension, uses model default if empty')
+parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN',
+                    help='Override mean pixel value of dataset')
+parser.add_argument('--std', type=float,  nargs='+', default=None, metavar='STD',
+                    help='Override std deviation of of dataset')
+parser.add_argument('--crop-pct', type=float, default=None, metavar='PCT',
+                    help='Override default crop pct of 0.875')
+parser.add_argument('--interpolation', default='', type=str, metavar='NAME',
+                    help='Image resize interpolation type (overrides model)')
+parser.add_argument('--tf-preprocessing', dest='tf_preprocessing', action='store_true',
+                    help='use tensorflow mnasnet preporcessing')
+parser.add_argument('--print-freq', '-p', default=10, type=int,
+                    metavar='N', help='print frequency (default: 10)')
+
+
+def main():
+    args = parser.parse_args()
+    args.gpu_id = 0
+    if args.c2_prefix:
+        args.c2_init = args.c2_prefix + '.init.pb'
+        args.c2_predict = args.c2_prefix + '.predict.pb'
+
+    model = model_helper.ModelHelper(name="validation_net", init_params=False)
+
+    # Bring in the init net from init_net.pb
+    init_net_proto = caffe2_pb2.NetDef()
+    with open(args.c2_init, "rb") as f:
+        init_net_proto.ParseFromString(f.read())
+    model.param_init_net = core.Net(init_net_proto)
+
+    # bring in the predict net from predict_net.pb
+    predict_net_proto = caffe2_pb2.NetDef()
+    with open(args.c2_predict, "rb") as f:
+        predict_net_proto.ParseFromString(f.read())
+    model.net = core.Net(predict_net_proto)
+
+    data_config = resolve_data_config(None, args)
+    loader = create_loader(
+        Dataset(args.data, load_bytes=args.tf_preprocessing),
+        input_size=data_config['input_size'],
+        batch_size=args.batch_size,
+        use_prefetcher=False,
+        interpolation=data_config['interpolation'],
+        mean=data_config['mean'],
+        std=data_config['std'],
+        num_workers=args.workers,
+        crop_pct=data_config['crop_pct'],
+        tensorflow_preprocessing=args.tf_preprocessing)
+
+    # this is so obvious, wonderful interface </sarcasm>
+    input_blob = model.net.external_inputs[0]
+    output_blob = model.net.external_outputs[0]
+
+    if True:
+        device_opts = None
+    else:
+        # CUDA is crashing, no idea why, awesome error message, give it a try for kicks
+        device_opts = core.DeviceOption(caffe2_pb2.PROTO_CUDA, args.gpu_id)
+        model.net.RunAllOnGPU(gpu_id=args.gpu_id, use_cudnn=True)
+        model.param_init_net.RunAllOnGPU(gpu_id=args.gpu_id, use_cudnn=True)
+
+    model.param_init_net.GaussianFill(
+        [], input_blob.GetUnscopedName(),
+        shape=(1,) + data_config['input_size'], mean=0.0, std=1.0)
+    workspace.RunNetOnce(model.param_init_net)
+    workspace.CreateNet(model.net, overwrite=True)
+
+    batch_time = AverageMeter()
+    top1 = AverageMeter()
+    top5 = AverageMeter()
+    end = time.time()
+    for i, (input, target) in enumerate(loader):
+        # run the net and return prediction
+        caffe2_in = input.data.numpy()
+        workspace.FeedBlob(input_blob, caffe2_in, device_opts)
+        workspace.RunNet(model.net, num_iter=1)
+        output = workspace.FetchBlob(output_blob)
+
+        # measure accuracy and record loss
+        prec1, prec5 = accuracy_np(output.data, target.numpy())
+        top1.update(prec1.item(), input.size(0))
+        top5.update(prec5.item(), input.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if i % args.print_freq == 0:
+            print('Test: [{0}/{1}]\t'
+                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f}, {rate_avg:.3f}/s, {ms_avg:.3f} ms/sample) \t'
+                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
+                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
+                i, len(loader), batch_time=batch_time, rate_avg=input.size(0) / batch_time.avg,
+                ms_avg=100 * batch_time.avg / input.size(0), top1=top1, top5=top5))
+
+    print(' * Prec@1 {top1.avg:.3f} ({top1a:.3f}) Prec@5 {top5.avg:.3f} ({top5a:.3f})'.format(
+        top1=top1, top1a=100-top1.avg, top5=top5, top5a=100.-top5.avg))
+
+
+def accuracy_np(output, target):
+    max_indices = np.argsort(output, axis=1)[:, ::-1]
+    top5 = 100 * np.equal(max_indices[:, :5], target[:, np.newaxis]).sum(axis=1).mean()
+    top1 = 100 * np.equal(max_indices[:, 0], target).mean()
+    return top1, top5
+
+
+if __name__ == '__main__':
+    main()

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/__init__.py

@@ -0,0 +1,5 @@
+from .gen_efficientnet import *
+from .mobilenetv3 import *
+from .model_factory import create_model
+from .config import is_exportable, is_scriptable, set_exportable, set_scriptable
+from .activations import *

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/activations/__init__.py

@@ -0,0 +1,137 @@
+from geffnet import config
+from geffnet.activations.activations_me import *
+from geffnet.activations.activations_jit import *
+from geffnet.activations.activations import *
+import torch
+
+_has_silu = 'silu' in dir(torch.nn.functional)
+
+_ACT_FN_DEFAULT = dict(
+    silu=F.silu if _has_silu else swish,
+    swish=F.silu if _has_silu else swish,
+    mish=mish,
+    relu=F.relu,
+    relu6=F.relu6,
+    sigmoid=sigmoid,
+    tanh=tanh,
+    hard_sigmoid=hard_sigmoid,
+    hard_swish=hard_swish,
+)
+
+_ACT_FN_JIT = dict(
+    silu=F.silu if _has_silu else swish_jit,
+    swish=F.silu if _has_silu else swish_jit,
+    mish=mish_jit,
+)
+
+_ACT_FN_ME = dict(
+    silu=F.silu if _has_silu else swish_me,
+    swish=F.silu if _has_silu else swish_me,
+    mish=mish_me,
+    hard_swish=hard_swish_me,
+    hard_sigmoid_jit=hard_sigmoid_me,
+)
+
+_ACT_LAYER_DEFAULT = dict(
+    silu=nn.SiLU if _has_silu else Swish,
+    swish=nn.SiLU if _has_silu else Swish,
+    mish=Mish,
+    relu=nn.ReLU,
+    relu6=nn.ReLU6,
+    sigmoid=Sigmoid,
+    tanh=Tanh,
+    hard_sigmoid=HardSigmoid,
+    hard_swish=HardSwish,
+)
+
+_ACT_LAYER_JIT = dict(
+    silu=nn.SiLU if _has_silu else SwishJit,
+    swish=nn.SiLU if _has_silu else SwishJit,
+    mish=MishJit,
+)
+
+_ACT_LAYER_ME = dict(
+    silu=nn.SiLU if _has_silu else SwishMe,
+    swish=nn.SiLU if _has_silu else SwishMe,
+    mish=MishMe,
+    hard_swish=HardSwishMe,
+    hard_sigmoid=HardSigmoidMe
+)
+
+_OVERRIDE_FN = dict()
+_OVERRIDE_LAYER = dict()
+
+
+def add_override_act_fn(name, fn):
+    global _OVERRIDE_FN
+    _OVERRIDE_FN[name] = fn
+
+
+def update_override_act_fn(overrides):
+    assert isinstance(overrides, dict)
+    global _OVERRIDE_FN
+    _OVERRIDE_FN.update(overrides)
+
+
+def clear_override_act_fn():
+    global _OVERRIDE_FN
+    _OVERRIDE_FN = dict()
+
+
+def add_override_act_layer(name, fn):
+    _OVERRIDE_LAYER[name] = fn
+
+
+def update_override_act_layer(overrides):
+    assert isinstance(overrides, dict)
+    global _OVERRIDE_LAYER
+    _OVERRIDE_LAYER.update(overrides)
+
+
+def clear_override_act_layer():
+    global _OVERRIDE_LAYER
+    _OVERRIDE_LAYER = dict()
+
+
+def get_act_fn(name='relu'):
+    """ Activation Function Factory
+    Fetching activation fns by name with this function allows export or torch script friendly
+    functions to be returned dynamically based on current config.
+    """
+    if name in _OVERRIDE_FN:
+        return _OVERRIDE_FN[name]
+    use_me = not (config.is_exportable() or config.is_scriptable() or config.is_no_jit())
+    if use_me and name in _ACT_FN_ME:
+        # If not exporting or scripting the model, first look for a memory optimized version
+        # activation with custom autograd, then fallback to jit scripted, then a Python or Torch builtin
+        return _ACT_FN_ME[name]
+    if config.is_exportable() and name in ('silu', 'swish'):
+        # FIXME PyTorch SiLU doesn't ONNX export, this is a temp hack
+        return swish
+    use_jit = not (config.is_exportable() or config.is_no_jit())
+    # NOTE: export tracing should work with jit scripted components, but I keep running into issues
+    if use_jit and name in _ACT_FN_JIT:  # jit scripted models should be okay for export/scripting
+        return _ACT_FN_JIT[name]
+    return _ACT_FN_DEFAULT[name]
+
+
+def get_act_layer(name='relu'):
+    """ Activation Layer Factory
+    Fetching activation layers by name with this function allows export or torch script friendly
+    functions to be returned dynamically based on current config.
+    """
+    if name in _OVERRIDE_LAYER:
+        return _OVERRIDE_LAYER[name]
+    use_me = not (config.is_exportable() or config.is_scriptable() or config.is_no_jit())
+    if use_me and name in _ACT_LAYER_ME:
+        return _ACT_LAYER_ME[name]
+    if config.is_exportable() and name in ('silu', 'swish'):
+        # FIXME PyTorch SiLU doesn't ONNX export, this is a temp hack
+        return Swish
+    use_jit = not (config.is_exportable() or config.is_no_jit())
+    # NOTE: export tracing should work with jit scripted components, but I keep running into issues
+    if use_jit and name in _ACT_FN_JIT:  # jit scripted models should be okay for export/scripting
+        return _ACT_LAYER_JIT[name]
+    return _ACT_LAYER_DEFAULT[name]
+
+

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/activations/activations.py

@@ -0,0 +1,102 @@
+""" Activations
+
+A collection of activations fn and modules with a common interface so that they can
+easily be swapped. All have an `inplace` arg even if not used.
+
+Copyright 2020 Ross Wightman
+"""
+from torch import nn as nn
+from torch.nn import functional as F
+
+
+def swish(x, inplace: bool = False):
+    """Swish - Described originally as SiLU (https://arxiv.org/abs/1702.03118v3)
+    and also as Swish (https://arxiv.org/abs/1710.05941).
+
+    TODO Rename to SiLU with addition to PyTorch
+    """
+    return x.mul_(x.sigmoid()) if inplace else x.mul(x.sigmoid())
+
+
+class Swish(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(Swish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return swish(x, self.inplace)
+
+
+def mish(x, inplace: bool = False):
+    """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    """
+    return x.mul(F.softplus(x).tanh())
+
+
+class Mish(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(Mish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return mish(x, self.inplace)
+
+
+def sigmoid(x, inplace: bool = False):
+    return x.sigmoid_() if inplace else x.sigmoid()
+
+
+# PyTorch has this, but not with a consistent inplace argmument interface
+class Sigmoid(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(Sigmoid, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return x.sigmoid_() if self.inplace else x.sigmoid()
+
+
+def tanh(x, inplace: bool = False):
+    return x.tanh_() if inplace else x.tanh()
+
+
+# PyTorch has this, but not with a consistent inplace argmument interface
+class Tanh(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(Tanh, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return x.tanh_() if self.inplace else x.tanh()
+
+
+def hard_swish(x, inplace: bool = False):
+    inner = F.relu6(x + 3.).div_(6.)
+    return x.mul_(inner) if inplace else x.mul(inner)
+
+
+class HardSwish(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSwish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return hard_swish(x, self.inplace)
+
+
+def hard_sigmoid(x, inplace: bool = False):
+    if inplace:
+        return x.add_(3.).clamp_(0., 6.).div_(6.)
+    else:
+        return F.relu6(x + 3.) / 6.
+
+
+class HardSigmoid(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSigmoid, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return hard_sigmoid(x, self.inplace)
+
+

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/activations/activations_jit.py

@@ -0,0 +1,79 @@
+""" Activations (jit)
+
+A collection of jit-scripted activations fn and modules with a common interface so that they can
+easily be swapped. All have an `inplace` arg even if not used.
+
+All jit scripted activations are lacking in-place variations on purpose, scripted kernel fusion does not
+currently work across in-place op boundaries, thus performance is equal to or less than the non-scripted
+versions if they contain in-place ops.
+
+Copyright 2020 Ross Wightman
+"""
+
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+__all__ = ['swish_jit', 'SwishJit', 'mish_jit', 'MishJit',
+           'hard_sigmoid_jit', 'HardSigmoidJit', 'hard_swish_jit', 'HardSwishJit']
+
+
+@torch.jit.script
+def swish_jit(x, inplace: bool = False):
+    """Swish - Described originally as SiLU (https://arxiv.org/abs/1702.03118v3)
+    and also as Swish (https://arxiv.org/abs/1710.05941).
+
+    TODO Rename to SiLU with addition to PyTorch
+    """
+    return x.mul(x.sigmoid())
+
+
+@torch.jit.script
+def mish_jit(x, _inplace: bool = False):
+    """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    """
+    return x.mul(F.softplus(x).tanh())
+
+
+class SwishJit(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(SwishJit, self).__init__()
+
+    def forward(self, x):
+        return swish_jit(x)
+
+
+class MishJit(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(MishJit, self).__init__()
+
+    def forward(self, x):
+        return mish_jit(x)
+
+
+@torch.jit.script
+def hard_sigmoid_jit(x, inplace: bool = False):
+    # return F.relu6(x + 3.) / 6.
+    return (x + 3).clamp(min=0, max=6).div(6.)  # clamp seems ever so slightly faster?
+
+
+class HardSigmoidJit(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSigmoidJit, self).__init__()
+
+    def forward(self, x):
+        return hard_sigmoid_jit(x)
+
+
+@torch.jit.script
+def hard_swish_jit(x, inplace: bool = False):
+    # return x * (F.relu6(x + 3.) / 6)
+    return x * (x + 3).clamp(min=0, max=6).div(6.)  # clamp seems ever so slightly faster?
+
+
+class HardSwishJit(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSwishJit, self).__init__()
+
+    def forward(self, x):
+        return hard_swish_jit(x)

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/activations/activations_me.py

@@ -0,0 +1,174 @@
+""" Activations (memory-efficient w/ custom autograd)
+
+A collection of activations fn and modules with a common interface so that they can
+easily be swapped. All have an `inplace` arg even if not used.
+
+These activations are not compatible with jit scripting or ONNX export of the model, please use either
+the JIT or basic versions of the activations.
+
+Copyright 2020 Ross Wightman
+"""
+
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+
+__all__ = ['swish_me', 'SwishMe', 'mish_me', 'MishMe',
+           'hard_sigmoid_me', 'HardSigmoidMe', 'hard_swish_me', 'HardSwishMe']
+
+
+@torch.jit.script
+def swish_jit_fwd(x):
+    return x.mul(torch.sigmoid(x))
+
+
+@torch.jit.script
+def swish_jit_bwd(x, grad_output):
+    x_sigmoid = torch.sigmoid(x)
+    return grad_output * (x_sigmoid * (1 + x * (1 - x_sigmoid)))
+
+
+class SwishJitAutoFn(torch.autograd.Function):
+    """ torch.jit.script optimised Swish w/ memory-efficient checkpoint
+    Inspired by conversation btw Jeremy Howard & Adam Pazske
+    https://twitter.com/jeremyphoward/status/1188251041835315200
+
+    Swish - Described originally as SiLU (https://arxiv.org/abs/1702.03118v3)
+    and also as Swish (https://arxiv.org/abs/1710.05941).
+
+    TODO Rename to SiLU with addition to PyTorch
+    """
+
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return swish_jit_fwd(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x = ctx.saved_tensors[0]
+        return swish_jit_bwd(x, grad_output)
+
+
+def swish_me(x, inplace=False):
+    return SwishJitAutoFn.apply(x)
+
+
+class SwishMe(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(SwishMe, self).__init__()
+
+    def forward(self, x):
+        return SwishJitAutoFn.apply(x)
+
+
+@torch.jit.script
+def mish_jit_fwd(x):
+    return x.mul(torch.tanh(F.softplus(x)))
+
+
+@torch.jit.script
+def mish_jit_bwd(x, grad_output):
+    x_sigmoid = torch.sigmoid(x)
+    x_tanh_sp = F.softplus(x).tanh()
+    return grad_output.mul(x_tanh_sp + x * x_sigmoid * (1 - x_tanh_sp * x_tanh_sp))
+
+
+class MishJitAutoFn(torch.autograd.Function):
+    """ Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    A memory efficient, jit scripted variant of Mish
+    """
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return mish_jit_fwd(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x = ctx.saved_tensors[0]
+        return mish_jit_bwd(x, grad_output)
+
+
+def mish_me(x, inplace=False):
+    return MishJitAutoFn.apply(x)
+
+
+class MishMe(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(MishMe, self).__init__()
+
+    def forward(self, x):
+        return MishJitAutoFn.apply(x)
+
+
+@torch.jit.script
+def hard_sigmoid_jit_fwd(x, inplace: bool = False):
+    return (x + 3).clamp(min=0, max=6).div(6.)
+
+
+@torch.jit.script
+def hard_sigmoid_jit_bwd(x, grad_output):
+    m = torch.ones_like(x) * ((x >= -3.) & (x <= 3.)) / 6.
+    return grad_output * m
+
+
+class HardSigmoidJitAutoFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return hard_sigmoid_jit_fwd(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x = ctx.saved_tensors[0]
+        return hard_sigmoid_jit_bwd(x, grad_output)
+
+
+def hard_sigmoid_me(x, inplace: bool = False):
+    return HardSigmoidJitAutoFn.apply(x)
+
+
+class HardSigmoidMe(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSigmoidMe, self).__init__()
+
+    def forward(self, x):
+        return HardSigmoidJitAutoFn.apply(x)
+
+
+@torch.jit.script
+def hard_swish_jit_fwd(x):
+    return x * (x + 3).clamp(min=0, max=6).div(6.)
+
+
+@torch.jit.script
+def hard_swish_jit_bwd(x, grad_output):
+    m = torch.ones_like(x) * (x >= 3.)
+    m = torch.where((x >= -3.) & (x <= 3.),  x / 3. + .5, m)
+    return grad_output * m
+
+
+class HardSwishJitAutoFn(torch.autograd.Function):
+    """A memory efficient, jit-scripted HardSwish activation"""
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return hard_swish_jit_fwd(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x = ctx.saved_tensors[0]
+        return hard_swish_jit_bwd(x, grad_output)
+
+
+def hard_swish_me(x, inplace=False):
+    return HardSwishJitAutoFn.apply(x)
+
+
+class HardSwishMe(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSwishMe, self).__init__()
+
+    def forward(self, x):
+        return HardSwishJitAutoFn.apply(x)

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/config.py

@@ -0,0 +1,123 @@
+""" Global layer config state
+"""
+from typing import Any, Optional
+
+__all__ = [
+    'is_exportable', 'is_scriptable', 'is_no_jit', 'layer_config_kwargs',
+    'set_exportable', 'set_scriptable', 'set_no_jit', 'set_layer_config'
+]
+
+# Set to True if prefer to have layers with no jit optimization (includes activations)
+_NO_JIT = False
+
+# Set to True if prefer to have activation layers with no jit optimization
+# NOTE not currently used as no difference between no_jit and no_activation jit as only layers obeying
+# the jit flags so far are activations. This will change as more layers are updated and/or added.
+_NO_ACTIVATION_JIT = False
+
+# Set to True if exporting a model with Same padding via ONNX
+_EXPORTABLE = False
+
+# Set to True if wanting to use torch.jit.script on a model
+_SCRIPTABLE = False
+
+
+def is_no_jit():
+    return _NO_JIT
+
+
+class set_no_jit:
+    def __init__(self, mode: bool) -> None:
+        global _NO_JIT
+        self.prev = _NO_JIT
+        _NO_JIT = mode
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(self, *args: Any) -> bool:
+        global _NO_JIT
+        _NO_JIT = self.prev
+        return False
+
+
+def is_exportable():
+    return _EXPORTABLE
+
+
+class set_exportable:
+    def __init__(self, mode: bool) -> None:
+        global _EXPORTABLE
+        self.prev = _EXPORTABLE
+        _EXPORTABLE = mode
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(self, *args: Any) -> bool:
+        global _EXPORTABLE
+        _EXPORTABLE = self.prev
+        return False
+
+
+def is_scriptable():
+    return _SCRIPTABLE
+
+
+class set_scriptable:
+    def __init__(self, mode: bool) -> None:
+        global _SCRIPTABLE
+        self.prev = _SCRIPTABLE
+        _SCRIPTABLE = mode
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(self, *args: Any) -> bool:
+        global _SCRIPTABLE
+        _SCRIPTABLE = self.prev
+        return False
+
+
+class set_layer_config:
+    """ Layer config context manager that allows setting all layer config flags at once.
+    If a flag arg is None, it will not change the current value.
+    """
+    def __init__(
+            self,
+            scriptable: Optional[bool] = None,
+            exportable: Optional[bool] = None,
+            no_jit: Optional[bool] = None,
+            no_activation_jit: Optional[bool] = None):
+        global _SCRIPTABLE
+        global _EXPORTABLE
+        global _NO_JIT
+        global _NO_ACTIVATION_JIT
+        self.prev = _SCRIPTABLE, _EXPORTABLE, _NO_JIT, _NO_ACTIVATION_JIT
+        if scriptable is not None:
+            _SCRIPTABLE = scriptable
+        if exportable is not None:
+            _EXPORTABLE = exportable
+        if no_jit is not None:
+            _NO_JIT = no_jit
+        if no_activation_jit is not None:
+            _NO_ACTIVATION_JIT = no_activation_jit
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(self, *args: Any) -> bool:
+        global _SCRIPTABLE
+        global _EXPORTABLE
+        global _NO_JIT
+        global _NO_ACTIVATION_JIT
+        _SCRIPTABLE, _EXPORTABLE, _NO_JIT, _NO_ACTIVATION_JIT = self.prev
+        return False
+
+
+def layer_config_kwargs(kwargs):
+    """ Consume config kwargs and return contextmgr obj """
+    return set_layer_config(
+        scriptable=kwargs.pop('scriptable', None),
+        exportable=kwargs.pop('exportable', None),
+        no_jit=kwargs.pop('no_jit', None))

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/conv2d_layers.py

@@ -0,0 +1,304 @@
+""" Conv2D w/ SAME padding, CondConv, MixedConv
+
+A collection of conv layers and padding helpers needed by EfficientNet, MixNet, and
+MobileNetV3 models that maintain weight compatibility with original Tensorflow models.
+
+Copyright 2020 Ross Wightman
+"""
+import collections.abc
+import math
+from functools import partial
+from itertools import repeat
+from typing import Tuple, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import is_exportable, is_scriptable
+
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+    return parse
+
+
+_single = _ntuple(1)
+_pair = _ntuple(2)
+_triple = _ntuple(3)
+_quadruple = _ntuple(4)
+
+
+def _is_static_pad(kernel_size, stride=1, dilation=1, **_):
+    return stride == 1 and (dilation * (kernel_size - 1)) % 2 == 0
+
+
+def _get_padding(kernel_size, stride=1, dilation=1, **_):
+    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
+    return padding
+
+
+def _calc_same_pad(i: int, k: int, s: int, d: int):
+    return max((-(i // -s) - 1) * s + (k - 1) * d + 1 - i, 0)
+
+
+def _same_pad_arg(input_size, kernel_size, stride, dilation):
+    ih, iw = input_size
+    kh, kw = kernel_size
+    pad_h = _calc_same_pad(ih, kh, stride[0], dilation[0])
+    pad_w = _calc_same_pad(iw, kw, stride[1], dilation[1])
+    return [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]
+
+
+def _split_channels(num_chan, num_groups):
+    split = [num_chan // num_groups for _ in range(num_groups)]
+    split[0] += num_chan - sum(split)
+    return split
+
+
+def conv2d_same(
+        x, weight: torch.Tensor, bias: Optional[torch.Tensor] = None, stride: Tuple[int, int] = (1, 1),
+        padding: Tuple[int, int] = (0, 0), dilation: Tuple[int, int] = (1, 1), groups: int = 1):
+    ih, iw = x.size()[-2:]
+    kh, kw = weight.size()[-2:]
+    pad_h = _calc_same_pad(ih, kh, stride[0], dilation[0])
+    pad_w = _calc_same_pad(iw, kw, stride[1], dilation[1])
+    x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2])
+    return F.conv2d(x, weight, bias, stride, (0, 0), dilation, groups)
+
+
+class Conv2dSame(nn.Conv2d):
+    """ Tensorflow like 'SAME' convolution wrapper for 2D convolutions
+    """
+
+    # pylint: disable=unused-argument
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                 padding=0, dilation=1, groups=1, bias=True):
+        super(Conv2dSame, self).__init__(
+            in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias)
+
+    def forward(self, x):
+        return conv2d_same(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+
+class Conv2dSameExport(nn.Conv2d):
+    """ ONNX export friendly Tensorflow like 'SAME' convolution wrapper for 2D convolutions
+
+    NOTE: This does not currently work with torch.jit.script
+    """
+
+    # pylint: disable=unused-argument
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                 padding=0, dilation=1, groups=1, bias=True):
+        super(Conv2dSameExport, self).__init__(
+            in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias)
+        self.pad = None
+        self.pad_input_size = (0, 0)
+
+    def forward(self, x):
+        input_size = x.size()[-2:]
+        if self.pad is None:
+            pad_arg = _same_pad_arg(input_size, self.weight.size()[-2:], self.stride, self.dilation)
+            self.pad = nn.ZeroPad2d(pad_arg)
+            self.pad_input_size = input_size
+
+        if self.pad is not None:
+            x = self.pad(x)
+        return F.conv2d(
+            x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+
+def get_padding_value(padding, kernel_size, **kwargs):
+    dynamic = False
+    if isinstance(padding, str):
+        # for any string padding, the padding will be calculated for you, one of three ways
+        padding = padding.lower()
+        if padding == 'same':
+            # TF compatible 'SAME' padding, has a performance and GPU memory allocation impact
+            if _is_static_pad(kernel_size, **kwargs):
+                # static case, no extra overhead
+                padding = _get_padding(kernel_size, **kwargs)
+            else:
+                # dynamic padding
+                padding = 0
+                dynamic = True
+        elif padding == 'valid':
+            # 'VALID' padding, same as padding=0
+            padding = 0
+        else:
+            # Default to PyTorch style 'same'-ish symmetric padding
+            padding = _get_padding(kernel_size, **kwargs)
+    return padding, dynamic
+
+
+def create_conv2d_pad(in_chs, out_chs, kernel_size, **kwargs):
+    padding = kwargs.pop('padding', '')
+    kwargs.setdefault('bias', False)
+    padding, is_dynamic = get_padding_value(padding, kernel_size, **kwargs)
+    if is_dynamic:
+        if is_exportable():
+            assert not is_scriptable()
+            return Conv2dSameExport(in_chs, out_chs, kernel_size, **kwargs)
+        else:
+            return Conv2dSame(in_chs, out_chs, kernel_size, **kwargs)
+    else:
+        return nn.Conv2d(in_chs, out_chs, kernel_size, padding=padding, **kwargs)
+
+
+class MixedConv2d(nn.ModuleDict):
+    """ Mixed Grouped Convolution
+    Based on MDConv and GroupedConv in MixNet impl:
+      https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mixnet/custom_layers.py
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size=3,
+                 stride=1, padding='', dilation=1, depthwise=False, **kwargs):
+        super(MixedConv2d, self).__init__()
+
+        kernel_size = kernel_size if isinstance(kernel_size, list) else [kernel_size]
+        num_groups = len(kernel_size)
+        in_splits = _split_channels(in_channels, num_groups)
+        out_splits = _split_channels(out_channels, num_groups)
+        self.in_channels = sum(in_splits)
+        self.out_channels = sum(out_splits)
+        for idx, (k, in_ch, out_ch) in enumerate(zip(kernel_size, in_splits, out_splits)):
+            conv_groups = out_ch if depthwise else 1
+            self.add_module(
+                str(idx),
+                create_conv2d_pad(
+                    in_ch, out_ch, k, stride=stride,
+                    padding=padding, dilation=dilation, groups=conv_groups, **kwargs)
+            )
+        self.splits = in_splits
+
+    def forward(self, x):
+        x_split = torch.split(x, self.splits, 1)
+        x_out = [conv(x_split[i]) for i, conv in enumerate(self.values())]
+        x = torch.cat(x_out, 1)
+        return x
+
+
+def get_condconv_initializer(initializer, num_experts, expert_shape):
+    def condconv_initializer(weight):
+        """CondConv initializer function."""
+        num_params = np.prod(expert_shape)
+        if (len(weight.shape) != 2 or weight.shape[0] != num_experts or
+                weight.shape[1] != num_params):
+            raise (ValueError(
+                'CondConv variables must have shape [num_experts, num_params]'))
+        for i in range(num_experts):
+            initializer(weight[i].view(expert_shape))
+    return condconv_initializer
+
+
+class CondConv2d(nn.Module):
+    """ Conditional Convolution
+    Inspired by: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/condconv/condconv_layers.py
+
+    Grouped convolution hackery for parallel execution of the per-sample kernel filters inspired by this discussion:
+    https://github.com/pytorch/pytorch/issues/17983
+    """
+    __constants__ = ['bias', 'in_channels', 'out_channels', 'dynamic_padding']
+
+    def __init__(self, in_channels, out_channels, kernel_size=3,
+                 stride=1, padding='', dilation=1, groups=1, bias=False, num_experts=4):
+        super(CondConv2d, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = _pair(stride)
+        padding_val, is_padding_dynamic = get_padding_value(
+            padding, kernel_size, stride=stride, dilation=dilation)
+        self.dynamic_padding = is_padding_dynamic  # if in forward to work with torchscript
+        self.padding = _pair(padding_val)
+        self.dilation = _pair(dilation)
+        self.groups = groups
+        self.num_experts = num_experts
+
+        self.weight_shape = (self.out_channels, self.in_channels // self.groups) + self.kernel_size
+        weight_num_param = 1
+        for wd in self.weight_shape:
+            weight_num_param *= wd
+        self.weight = torch.nn.Parameter(torch.Tensor(self.num_experts, weight_num_param))
+
+        if bias:
+            self.bias_shape = (self.out_channels,)
+            self.bias = torch.nn.Parameter(torch.Tensor(self.num_experts, self.out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init_weight = get_condconv_initializer(
+            partial(nn.init.kaiming_uniform_, a=math.sqrt(5)), self.num_experts, self.weight_shape)
+        init_weight(self.weight)
+        if self.bias is not None:
+            fan_in = np.prod(self.weight_shape[1:])
+            bound = 1 / math.sqrt(fan_in)
+            init_bias = get_condconv_initializer(
+                partial(nn.init.uniform_, a=-bound, b=bound), self.num_experts, self.bias_shape)
+            init_bias(self.bias)
+
+    def forward(self, x, routing_weights):
+        B, C, H, W = x.shape
+        weight = torch.matmul(routing_weights, self.weight)
+        new_weight_shape = (B * self.out_channels, self.in_channels // self.groups) + self.kernel_size
+        weight = weight.view(new_weight_shape)
+        bias = None
+        if self.bias is not None:
+            bias = torch.matmul(routing_weights, self.bias)
+            bias = bias.view(B * self.out_channels)
+        # move batch elements with channels so each batch element can be efficiently convolved with separate kernel
+        x = x.view(1, B * C, H, W)
+        if self.dynamic_padding:
+            out = conv2d_same(
+                x, weight, bias, stride=self.stride, padding=self.padding,
+                dilation=self.dilation, groups=self.groups * B)
+        else:
+            out = F.conv2d(
+                x, weight, bias, stride=self.stride, padding=self.padding,
+                dilation=self.dilation, groups=self.groups * B)
+        out = out.permute([1, 0, 2, 3]).view(B, self.out_channels, out.shape[-2], out.shape[-1])
+
+        # Literal port (from TF definition)
+        # x = torch.split(x, 1, 0)
+        # weight = torch.split(weight, 1, 0)
+        # if self.bias is not None:
+        #     bias = torch.matmul(routing_weights, self.bias)
+        #     bias = torch.split(bias, 1, 0)
+        # else:
+        #     bias = [None] * B
+        # out = []
+        # for xi, wi, bi in zip(x, weight, bias):
+        #     wi = wi.view(*self.weight_shape)
+        #     if bi is not None:
+        #         bi = bi.view(*self.bias_shape)
+        #     out.append(self.conv_fn(
+        #         xi, wi, bi, stride=self.stride, padding=self.padding,
+        #         dilation=self.dilation, groups=self.groups))
+        # out = torch.cat(out, 0)
+        return out
+
+
+def select_conv2d(in_chs, out_chs, kernel_size, **kwargs):
+    assert 'groups' not in kwargs  # only use 'depthwise' bool arg
+    if isinstance(kernel_size, list):
+        assert 'num_experts' not in kwargs  # MixNet + CondConv combo not supported currently
+        # We're going to use only lists for defining the MixedConv2d kernel groups,
+        # ints, tuples, other iterables will continue to pass to normal conv and specify h, w.
+        m = MixedConv2d(in_chs, out_chs, kernel_size, **kwargs)
+    else:
+        depthwise = kwargs.pop('depthwise', False)
+        groups = out_chs if depthwise else 1
+        if 'num_experts' in kwargs and kwargs['num_experts'] > 0:
+            m = CondConv2d(in_chs, out_chs, kernel_size, groups=groups, **kwargs)
+        else:
+            m = create_conv2d_pad(in_chs, out_chs, kernel_size, groups=groups, **kwargs)
+    return m

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/efficientnet_builder.py

@@ -0,0 +1,683 @@
+""" EfficientNet / MobileNetV3 Blocks and Builder
+
+Copyright 2020 Ross Wightman
+"""
+import re
+from copy import deepcopy
+
+from .conv2d_layers import CondConv2d, get_condconv_initializer, math, partial, select_conv2d
+from geffnet.activations import F, get_act_layer, nn, sigmoid, torch
+
+__all__ = ['get_bn_args_tf', 'resolve_bn_args', 'resolve_se_args', 'resolve_act_layer', 'make_divisible',
+           'round_channels', 'drop_connect', 'SqueezeExcite', 'ConvBnAct', 'DepthwiseSeparableConv',
+           'InvertedResidual', 'CondConvResidual', 'EdgeResidual', 'EfficientNetBuilder', 'decode_arch_def',
+           'initialize_weight_default', 'initialize_weight_goog', 'BN_MOMENTUM_TF_DEFAULT', 'BN_EPS_TF_DEFAULT'
+]
+
+# Defaults used for Google/Tensorflow training of mobile networks /w RMSprop as per
+# papers and TF reference implementations. PT momentum equiv for TF decay is (1 - TF decay)
+# NOTE: momentum varies btw .99 and .9997 depending on source
+# .99 in official TF TPU impl
+# .9997 (/w .999 in search space) for paper
+#
+# PyTorch defaults are momentum = .1, eps = 1e-5
+#
+BN_MOMENTUM_TF_DEFAULT = 1 - 0.99
+BN_EPS_TF_DEFAULT = 1e-3
+_BN_ARGS_TF = dict(momentum=BN_MOMENTUM_TF_DEFAULT, eps=BN_EPS_TF_DEFAULT)
+
+
+def get_bn_args_tf():
+    return _BN_ARGS_TF.copy()
+
+
+def resolve_bn_args(kwargs):
+    bn_args = get_bn_args_tf() if kwargs.pop('bn_tf', False) else {}
+    bn_momentum = kwargs.pop('bn_momentum', None)
+    if bn_momentum is not None:
+        bn_args['momentum'] = bn_momentum
+    bn_eps = kwargs.pop('bn_eps', None)
+    if bn_eps is not None:
+        bn_args['eps'] = bn_eps
+    return bn_args
+
+
+_SE_ARGS_DEFAULT = dict(
+    gate_fn=sigmoid,
+    act_layer=None,  # None == use containing block's activation layer
+    reduce_mid=False,
+    divisor=1)
+
+
+def resolve_se_args(kwargs, in_chs, act_layer=None):
+    se_kwargs = kwargs.copy() if kwargs is not None else {}
+    # fill in args that aren't specified with the defaults
+    for k, v in _SE_ARGS_DEFAULT.items():
+        se_kwargs.setdefault(k, v)
+    # some models, like MobilNetV3, calculate SE reduction chs from the containing block's mid_ch instead of in_ch
+    if not se_kwargs.pop('reduce_mid'):
+        se_kwargs['reduced_base_chs'] = in_chs
+    # act_layer override, if it remains None, the containing block's act_layer will be used
+    if se_kwargs['act_layer'] is None:
+        assert act_layer is not None
+        se_kwargs['act_layer'] = act_layer
+    return se_kwargs
+
+
+def resolve_act_layer(kwargs, default='relu'):
+    act_layer = kwargs.pop('act_layer', default)
+    if isinstance(act_layer, str):
+        act_layer = get_act_layer(act_layer)
+    return act_layer
+
+
+def make_divisible(v: int, divisor: int = 8, min_value: int = None):
+    min_value = min_value or divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    if new_v < 0.9 * v:  # ensure round down does not go down by more than 10%.
+        new_v += divisor
+    return new_v
+
+
+def round_channels(channels, multiplier=1.0, divisor=8, channel_min=None):
+    """Round number of filters based on depth multiplier."""
+    if not multiplier:
+        return channels
+    channels *= multiplier
+    return make_divisible(channels, divisor, channel_min)
+
+
+def drop_connect(inputs, training: bool = False, drop_connect_rate: float = 0.):
+    """Apply drop connect."""
+    if not training:
+        return inputs
+
+    keep_prob = 1 - drop_connect_rate
+    random_tensor = keep_prob + torch.rand(
+        (inputs.size()[0], 1, 1, 1), dtype=inputs.dtype, device=inputs.device)
+    random_tensor.floor_()  # binarize
+    output = inputs.div(keep_prob) * random_tensor
+    return output
+
+
+class SqueezeExcite(nn.Module):
+
+    def __init__(self, in_chs, se_ratio=0.25, reduced_base_chs=None, act_layer=nn.ReLU, gate_fn=sigmoid, divisor=1):
+        super(SqueezeExcite, self).__init__()
+        reduced_chs = make_divisible((reduced_base_chs or in_chs) * se_ratio, divisor)
+        self.conv_reduce = nn.Conv2d(in_chs, reduced_chs, 1, bias=True)
+        self.act1 = act_layer(inplace=True)
+        self.conv_expand = nn.Conv2d(reduced_chs, in_chs, 1, bias=True)
+        self.gate_fn = gate_fn
+
+    def forward(self, x):
+        x_se = x.mean((2, 3), keepdim=True)
+        x_se = self.conv_reduce(x_se)
+        x_se = self.act1(x_se)
+        x_se = self.conv_expand(x_se)
+        x = x * self.gate_fn(x_se)
+        return x
+
+
+class ConvBnAct(nn.Module):
+    def __init__(self, in_chs, out_chs, kernel_size,
+                 stride=1, pad_type='', act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, norm_kwargs=None):
+        super(ConvBnAct, self).__init__()
+        assert stride in [1, 2]
+        norm_kwargs = norm_kwargs or {}
+        self.conv = select_conv2d(in_chs, out_chs, kernel_size, stride=stride, padding=pad_type)
+        self.bn1 = norm_layer(out_chs, **norm_kwargs)
+        self.act1 = act_layer(inplace=True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+        return x
+
+
+class DepthwiseSeparableConv(nn.Module):
+    """ DepthwiseSeparable block
+    Used for DS convs in MobileNet-V1 and in the place of IR blocks with an expansion
+    factor of 1.0. This is an alternative to having a IR with optional first pw conv.
+    """
+    def __init__(self, in_chs, out_chs, dw_kernel_size=3,
+                 stride=1, pad_type='', act_layer=nn.ReLU, noskip=False,
+                 pw_kernel_size=1, pw_act=False, se_ratio=0., se_kwargs=None,
+                 norm_layer=nn.BatchNorm2d, norm_kwargs=None, drop_connect_rate=0.):
+        super(DepthwiseSeparableConv, self).__init__()
+        assert stride in [1, 2]
+        norm_kwargs = norm_kwargs or {}
+        self.has_residual = (stride == 1 and in_chs == out_chs) and not noskip
+        self.drop_connect_rate = drop_connect_rate
+
+        self.conv_dw = select_conv2d(
+            in_chs, in_chs, dw_kernel_size, stride=stride, padding=pad_type, depthwise=True)
+        self.bn1 = norm_layer(in_chs, **norm_kwargs)
+        self.act1 = act_layer(inplace=True)
+
+        # Squeeze-and-excitation
+        if se_ratio is not None and se_ratio > 0.:
+            se_kwargs = resolve_se_args(se_kwargs, in_chs, act_layer)
+            self.se = SqueezeExcite(in_chs, se_ratio=se_ratio, **se_kwargs)
+        else:
+            self.se = nn.Identity()
+
+        self.conv_pw = select_conv2d(in_chs, out_chs, pw_kernel_size, padding=pad_type)
+        self.bn2 = norm_layer(out_chs, **norm_kwargs)
+        self.act2 = act_layer(inplace=True) if pw_act else nn.Identity()
+
+    def forward(self, x):
+        residual = x
+
+        x = self.conv_dw(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+
+        x = self.se(x)
+
+        x = self.conv_pw(x)
+        x = self.bn2(x)
+        x = self.act2(x)
+
+        if self.has_residual:
+            if self.drop_connect_rate > 0.:
+                x = drop_connect(x, self.training, self.drop_connect_rate)
+            x += residual
+        return x
+
+
+class InvertedResidual(nn.Module):
+    """ Inverted residual block w/ optional SE"""
+
+    def __init__(self, in_chs, out_chs, dw_kernel_size=3,
+                 stride=1, pad_type='', act_layer=nn.ReLU, noskip=False,
+                 exp_ratio=1.0, exp_kernel_size=1, pw_kernel_size=1,
+                 se_ratio=0., se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None,
+                 conv_kwargs=None, drop_connect_rate=0.):
+        super(InvertedResidual, self).__init__()
+        norm_kwargs = norm_kwargs or {}
+        conv_kwargs = conv_kwargs or {}
+        mid_chs: int = make_divisible(in_chs * exp_ratio)
+        self.has_residual = (in_chs == out_chs and stride == 1) and not noskip
+        self.drop_connect_rate = drop_connect_rate
+
+        # Point-wise expansion
+        self.conv_pw = select_conv2d(in_chs, mid_chs, exp_kernel_size, padding=pad_type, **conv_kwargs)
+        self.bn1 = norm_layer(mid_chs, **norm_kwargs)
+        self.act1 = act_layer(inplace=True)
+
+        # Depth-wise convolution
+        self.conv_dw = select_conv2d(
+            mid_chs, mid_chs, dw_kernel_size, stride=stride, padding=pad_type, depthwise=True, **conv_kwargs)
+        self.bn2 = norm_layer(mid_chs, **norm_kwargs)
+        self.act2 = act_layer(inplace=True)
+
+        # Squeeze-and-excitation
+        if se_ratio is not None and se_ratio > 0.:
+            se_kwargs = resolve_se_args(se_kwargs, in_chs, act_layer)
+            self.se = SqueezeExcite(mid_chs, se_ratio=se_ratio, **se_kwargs)
+        else:
+            self.se = nn.Identity()  # for jit.script compat
+
+        # Point-wise linear projection
+        self.conv_pwl = select_conv2d(mid_chs, out_chs, pw_kernel_size, padding=pad_type, **conv_kwargs)
+        self.bn3 = norm_layer(out_chs, **norm_kwargs)
+
+    def forward(self, x):
+        residual = x
+
+        # Point-wise expansion
+        x = self.conv_pw(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+
+        # Depth-wise convolution
+        x = self.conv_dw(x)
+        x = self.bn2(x)
+        x = self.act2(x)
+
+        # Squeeze-and-excitation
+        x = self.se(x)
+
+        # Point-wise linear projection
+        x = self.conv_pwl(x)
+        x = self.bn3(x)
+
+        if self.has_residual:
+            if self.drop_connect_rate > 0.:
+                x = drop_connect(x, self.training, self.drop_connect_rate)
+            x += residual
+        return x
+
+
+class CondConvResidual(InvertedResidual):
+    """ Inverted residual block w/ CondConv routing"""
+
+    def __init__(self, in_chs, out_chs, dw_kernel_size=3,
+                 stride=1, pad_type='', act_layer=nn.ReLU, noskip=False,
+                 exp_ratio=1.0, exp_kernel_size=1, pw_kernel_size=1,
+                 se_ratio=0., se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None,
+                 num_experts=0, drop_connect_rate=0.):
+
+        self.num_experts = num_experts
+        conv_kwargs = dict(num_experts=self.num_experts)
+
+        super(CondConvResidual, self).__init__(
+            in_chs, out_chs, dw_kernel_size=dw_kernel_size, stride=stride, pad_type=pad_type,
+            act_layer=act_layer, noskip=noskip, exp_ratio=exp_ratio, exp_kernel_size=exp_kernel_size,
+            pw_kernel_size=pw_kernel_size, se_ratio=se_ratio, se_kwargs=se_kwargs,
+            norm_layer=norm_layer, norm_kwargs=norm_kwargs, conv_kwargs=conv_kwargs,
+            drop_connect_rate=drop_connect_rate)
+
+        self.routing_fn = nn.Linear(in_chs, self.num_experts)
+
+    def forward(self, x):
+        residual = x
+
+        # CondConv routing
+        pooled_inputs = F.adaptive_avg_pool2d(x, 1).flatten(1)
+        routing_weights = torch.sigmoid(self.routing_fn(pooled_inputs))
+
+        # Point-wise expansion
+        x = self.conv_pw(x, routing_weights)
+        x = self.bn1(x)
+        x = self.act1(x)
+
+        # Depth-wise convolution
+        x = self.conv_dw(x, routing_weights)
+        x = self.bn2(x)
+        x = self.act2(x)
+
+        # Squeeze-and-excitation
+        x = self.se(x)
+
+        # Point-wise linear projection
+        x = self.conv_pwl(x, routing_weights)
+        x = self.bn3(x)
+
+        if self.has_residual:
+            if self.drop_connect_rate > 0.:
+                x = drop_connect(x, self.training, self.drop_connect_rate)
+            x += residual
+        return x
+
+
+class EdgeResidual(nn.Module):
+    """ EdgeTPU Residual block with expansion convolution followed by pointwise-linear w/ stride"""
+
+    def __init__(self, in_chs, out_chs, exp_kernel_size=3, exp_ratio=1.0, fake_in_chs=0,
+                 stride=1, pad_type='', act_layer=nn.ReLU, noskip=False, pw_kernel_size=1,
+                 se_ratio=0., se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None, drop_connect_rate=0.):
+        super(EdgeResidual, self).__init__()
+        norm_kwargs = norm_kwargs or {}
+        mid_chs = make_divisible(fake_in_chs * exp_ratio) if fake_in_chs > 0 else make_divisible(in_chs * exp_ratio)
+        self.has_residual = (in_chs == out_chs and stride == 1) and not noskip
+        self.drop_connect_rate = drop_connect_rate
+
+        # Expansion convolution
+        self.conv_exp = select_conv2d(in_chs, mid_chs, exp_kernel_size, padding=pad_type)
+        self.bn1 = norm_layer(mid_chs, **norm_kwargs)
+        self.act1 = act_layer(inplace=True)
+
+        # Squeeze-and-excitation
+        if se_ratio is not None and se_ratio > 0.:
+            se_kwargs = resolve_se_args(se_kwargs, in_chs, act_layer)
+            self.se = SqueezeExcite(mid_chs, se_ratio=se_ratio, **se_kwargs)
+        else:
+            self.se = nn.Identity()
+
+        # Point-wise linear projection
+        self.conv_pwl = select_conv2d(mid_chs, out_chs, pw_kernel_size, stride=stride, padding=pad_type)
+        self.bn2 = nn.BatchNorm2d(out_chs, **norm_kwargs)
+
+    def forward(self, x):
+        residual = x
+
+        # Expansion convolution
+        x = self.conv_exp(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+
+        # Squeeze-and-excitation
+        x = self.se(x)
+
+        # Point-wise linear projection
+        x = self.conv_pwl(x)
+        x = self.bn2(x)
+
+        if self.has_residual:
+            if self.drop_connect_rate > 0.:
+                x = drop_connect(x, self.training, self.drop_connect_rate)
+            x += residual
+
+        return x
+
+
+class EfficientNetBuilder:
+    """ Build Trunk Blocks for Efficient/Mobile Networks
+
+    This ended up being somewhat of a cross between
+    https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mnasnet_models.py
+    and
+    https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/maskrcnn_benchmark/modeling/backbone/fbnet_builder.py
+
+    """
+
+    def __init__(self, channel_multiplier=1.0, channel_divisor=8, channel_min=None,
+                 pad_type='', act_layer=None, se_kwargs=None,
+                 norm_layer=nn.BatchNorm2d, norm_kwargs=None, drop_connect_rate=0.):
+        self.channel_multiplier = channel_multiplier
+        self.channel_divisor = channel_divisor
+        self.channel_min = channel_min
+        self.pad_type = pad_type
+        self.act_layer = act_layer
+        self.se_kwargs = se_kwargs
+        self.norm_layer = norm_layer
+        self.norm_kwargs = norm_kwargs
+        self.drop_connect_rate = drop_connect_rate
+
+        # updated during build
+        self.in_chs = None
+        self.block_idx = 0
+        self.block_count = 0
+
+    def _round_channels(self, chs):
+        return round_channels(chs, self.channel_multiplier, self.channel_divisor, self.channel_min)
+
+    def _make_block(self, ba):
+        bt = ba.pop('block_type')
+        ba['in_chs'] = self.in_chs
+        ba['out_chs'] = self._round_channels(ba['out_chs'])
+        if 'fake_in_chs' in ba and ba['fake_in_chs']:
+            # FIXME this is a hack to work around mismatch in origin impl input filters for EdgeTPU
+            ba['fake_in_chs'] = self._round_channels(ba['fake_in_chs'])
+        ba['norm_layer'] = self.norm_layer
+        ba['norm_kwargs'] = self.norm_kwargs
+        ba['pad_type'] = self.pad_type
+        # block act fn overrides the model default
+        ba['act_layer'] = ba['act_layer'] if ba['act_layer'] is not None else self.act_layer
+        assert ba['act_layer'] is not None
+        if bt == 'ir':
+            ba['drop_connect_rate'] = self.drop_connect_rate * self.block_idx / self.block_count
+            ba['se_kwargs'] = self.se_kwargs
+            if ba.get('num_experts', 0) > 0:
+                block = CondConvResidual(**ba)
+            else:
+                block = InvertedResidual(**ba)
+        elif bt == 'ds' or bt == 'dsa':
+            ba['drop_connect_rate'] = self.drop_connect_rate * self.block_idx / self.block_count
+            ba['se_kwargs'] = self.se_kwargs
+            block = DepthwiseSeparableConv(**ba)
+        elif bt == 'er':
+            ba['drop_connect_rate'] = self.drop_connect_rate * self.block_idx / self.block_count
+            ba['se_kwargs'] = self.se_kwargs
+            block = EdgeResidual(**ba)
+        elif bt == 'cn':
+            block = ConvBnAct(**ba)
+        else:
+            assert False, 'Uknkown block type (%s) while building model.' % bt
+        self.in_chs = ba['out_chs']  # update in_chs for arg of next block
+        return block
+
+    def _make_stack(self, stack_args):
+        blocks = []
+        # each stack (stage) contains a list of block arguments
+        for i, ba in enumerate(stack_args):
+            if i >= 1:
+                # only the first block in any stack can have a stride > 1
+                ba['stride'] = 1
+            block = self._make_block(ba)
+            blocks.append(block)
+            self.block_idx += 1  # incr global idx (across all stacks)
+        return nn.Sequential(*blocks)
+
+    def __call__(self, in_chs, block_args):
+        """ Build the blocks
+        Args:
+            in_chs: Number of input-channels passed to first block
+            block_args: A list of lists, outer list defines stages, inner
+                list contains strings defining block configuration(s)
+        Return:
+             List of block stacks (each stack wrapped in nn.Sequential)
+        """
+        self.in_chs = in_chs
+        self.block_count = sum([len(x) for x in block_args])
+        self.block_idx = 0
+        blocks = []
+        # outer list of block_args defines the stacks ('stages' by some conventions)
+        for stack_idx, stack in enumerate(block_args):
+            assert isinstance(stack, list)
+            stack = self._make_stack(stack)
+            blocks.append(stack)
+        return blocks
+
+
+def _parse_ksize(ss):
+    if ss.isdigit():
+        return int(ss)
+    else:
+        return [int(k) for k in ss.split('.')]
+
+
+def _decode_block_str(block_str):
+    """ Decode block definition string
+
+    Gets a list of block arg (dicts) through a string notation of arguments.
+    E.g. ir_r2_k3_s2_e1_i32_o16_se0.25_noskip
+
+    All args can exist in any order with the exception of the leading string which
+    is assumed to indicate the block type.
+
+    leading string - block type (
+      ir = InvertedResidual, ds = DepthwiseSep, dsa = DeptwhiseSep with pw act, cn = ConvBnAct)
+    r - number of repeat blocks,
+    k - kernel size,
+    s - strides (1-9),
+    e - expansion ratio,
+    c - output channels,
+    se - squeeze/excitation ratio
+    n - activation fn ('re', 'r6', 'hs', or 'sw')
+    Args:
+        block_str: a string representation of block arguments.
+    Returns:
+        A list of block args (dicts)
+    Raises:
+        ValueError: if the string def not properly specified (TODO)
+    """
+    assert isinstance(block_str, str)
+    ops = block_str.split('_')
+    block_type = ops[0]  # take the block type off the front
+    ops = ops[1:]
+    options = {}
+    noskip = False
+    for op in ops:
+        # string options being checked on individual basis, combine if they grow
+        if op == 'noskip':
+            noskip = True
+        elif op.startswith('n'):
+            # activation fn
+            key = op[0]
+            v = op[1:]
+            if v == 're':
+                value = get_act_layer('relu')
+            elif v == 'r6':
+                value = get_act_layer('relu6')
+            elif v == 'hs':
+                value = get_act_layer('hard_swish')
+            elif v == 'sw':
+                value = get_act_layer('swish')
+            else:
+                continue
+            options[key] = value
+        else:
+            # all numeric options
+            splits = re.split(r'(\d.*)', op)
+            if len(splits) >= 2:
+                key, value = splits[:2]
+                options[key] = value
+
+    # if act_layer is None, the model default (passed to model init) will be used
+    act_layer = options['n'] if 'n' in options else None
+    exp_kernel_size = _parse_ksize(options['a']) if 'a' in options else 1
+    pw_kernel_size = _parse_ksize(options['p']) if 'p' in options else 1
+    fake_in_chs = int(options['fc']) if 'fc' in options else 0  # FIXME hack to deal with in_chs issue in TPU def
+
+    num_repeat = int(options['r'])
+    # each type of block has different valid arguments, fill accordingly
+    if block_type == 'ir':
+        block_args = dict(
+            block_type=block_type,
+            dw_kernel_size=_parse_ksize(options['k']),
+            exp_kernel_size=exp_kernel_size,
+            pw_kernel_size=pw_kernel_size,
+            out_chs=int(options['c']),
+            exp_ratio=float(options['e']),
+            se_ratio=float(options['se']) if 'se' in options else None,
+            stride=int(options['s']),
+            act_layer=act_layer,
+            noskip=noskip,
+        )
+        if 'cc' in options:
+            block_args['num_experts'] = int(options['cc'])
+    elif block_type == 'ds' or block_type == 'dsa':
+        block_args = dict(
+            block_type=block_type,
+            dw_kernel_size=_parse_ksize(options['k']),
+            pw_kernel_size=pw_kernel_size,
+            out_chs=int(options['c']),
+            se_ratio=float(options['se']) if 'se' in options else None,
+            stride=int(options['s']),
+            act_layer=act_layer,
+            pw_act=block_type == 'dsa',
+            noskip=block_type == 'dsa' or noskip,
+        )
+    elif block_type == 'er':
+        block_args = dict(
+            block_type=block_type,
+            exp_kernel_size=_parse_ksize(options['k']),
+            pw_kernel_size=pw_kernel_size,
+            out_chs=int(options['c']),
+            exp_ratio=float(options['e']),
+            fake_in_chs=fake_in_chs,
+            se_ratio=float(options['se']) if 'se' in options else None,
+            stride=int(options['s']),
+            act_layer=act_layer,
+            noskip=noskip,
+        )
+    elif block_type == 'cn':
+        block_args = dict(
+            block_type=block_type,
+            kernel_size=int(options['k']),
+            out_chs=int(options['c']),
+            stride=int(options['s']),
+            act_layer=act_layer,
+        )
+    else:
+        assert False, 'Unknown block type (%s)' % block_type
+
+    return block_args, num_repeat
+
+
+def _scale_stage_depth(stack_args, repeats, depth_multiplier=1.0, depth_trunc='ceil'):
+    """ Per-stage depth scaling
+    Scales the block repeats in each stage. This depth scaling impl maintains
+    compatibility with the EfficientNet scaling method, while allowing sensible
+    scaling for other models that may have multiple block arg definitions in each stage.
+    """
+
+    # We scale the total repeat count for each stage, there may be multiple
+    # block arg defs per stage so we need to sum.
+    num_repeat = sum(repeats)
+    if depth_trunc == 'round':
+        # Truncating to int by rounding allows stages with few repeats to remain
+        # proportionally smaller for longer. This is a good choice when stage definitions
+        # include single repeat stages that we'd prefer to keep that way as long as possible
+        num_repeat_scaled = max(1, round(num_repeat * depth_multiplier))
+    else:
+        # The default for EfficientNet truncates repeats to int via 'ceil'.
+        # Any multiplier > 1.0 will result in an increased depth for every stage.
+        num_repeat_scaled = int(math.ceil(num_repeat * depth_multiplier))
+
+    # Proportionally distribute repeat count scaling to each block definition in the stage.
+    # Allocation is done in reverse as it results in the first block being less likely to be scaled.
+    # The first block makes less sense to repeat in most of the arch definitions.
+    repeats_scaled = []
+    for r in repeats[::-1]:
+        rs = max(1, round((r / num_repeat * num_repeat_scaled)))
+        repeats_scaled.append(rs)
+        num_repeat -= r
+        num_repeat_scaled -= rs
+    repeats_scaled = repeats_scaled[::-1]
+
+    # Apply the calculated scaling to each block arg in the stage
+    sa_scaled = []
+    for ba, rep in zip(stack_args, repeats_scaled):
+        sa_scaled.extend([deepcopy(ba) for _ in range(rep)])
+    return sa_scaled
+
+
+def decode_arch_def(arch_def, depth_multiplier=1.0, depth_trunc='ceil', experts_multiplier=1, fix_first_last=False):
+    arch_args = []
+    for stack_idx, block_strings in enumerate(arch_def):
+        assert isinstance(block_strings, list)
+        stack_args = []
+        repeats = []
+        for block_str in block_strings:
+            assert isinstance(block_str, str)
+            ba, rep = _decode_block_str(block_str)
+            if ba.get('num_experts', 0) > 0 and experts_multiplier > 1:
+                ba['num_experts'] *= experts_multiplier
+            stack_args.append(ba)
+            repeats.append(rep)
+        if fix_first_last and (stack_idx == 0 or stack_idx == len(arch_def) - 1):
+            arch_args.append(_scale_stage_depth(stack_args, repeats, 1.0, depth_trunc))
+        else:
+            arch_args.append(_scale_stage_depth(stack_args, repeats, depth_multiplier, depth_trunc))
+    return arch_args
+
+
+def initialize_weight_goog(m, n='', fix_group_fanout=True):
+    # weight init as per Tensorflow Official impl
+    # https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mnasnet_model.py
+    if isinstance(m, CondConv2d):
+        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+        if fix_group_fanout:
+            fan_out //= m.groups
+        init_weight_fn = get_condconv_initializer(
+            lambda w: w.data.normal_(0, math.sqrt(2.0 / fan_out)), m.num_experts, m.weight_shape)
+        init_weight_fn(m.weight)
+        if m.bias is not None:
+            m.bias.data.zero_()
+    elif isinstance(m, nn.Conv2d):
+        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+        if fix_group_fanout:
+            fan_out //= m.groups
+        m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+        if m.bias is not None:
+            m.bias.data.zero_()
+    elif isinstance(m, nn.BatchNorm2d):
+        m.weight.data.fill_(1.0)
+        m.bias.data.zero_()
+    elif isinstance(m, nn.Linear):
+        fan_out = m.weight.size(0)  # fan-out
+        fan_in = 0
+        if 'routing_fn' in n:
+            fan_in = m.weight.size(1)
+        init_range = 1.0 / math.sqrt(fan_in + fan_out)
+        m.weight.data.uniform_(-init_range, init_range)
+        m.bias.data.zero_()
+
+
+def initialize_weight_default(m, n=''):
+    if isinstance(m, CondConv2d):
+        init_fn = get_condconv_initializer(partial(
+            nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'), m.num_experts, m.weight_shape)
+        init_fn(m.weight)
+    elif isinstance(m, nn.Conv2d):
+        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+    elif isinstance(m, nn.BatchNorm2d):
+        m.weight.data.fill_(1.0)
+        m.bias.data.zero_()
+    elif isinstance(m, nn.Linear):
+        nn.init.kaiming_uniform_(m.weight, mode='fan_in', nonlinearity='linear')

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/helpers.py

@@ -0,0 +1,71 @@
+""" Checkpoint loading / state_dict helpers
+Copyright 2020 Ross Wightman
+"""
+import torch
+import os
+from collections import OrderedDict
+try:
+    from torch.hub import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+
+def load_checkpoint(model, checkpoint_path):
+    if checkpoint_path and os.path.isfile(checkpoint_path):
+        print("=> Loading checkpoint '{}'".format(checkpoint_path))
+        checkpoint = torch.load(checkpoint_path)
+        if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+            new_state_dict = OrderedDict()
+            for k, v in checkpoint['state_dict'].items():
+                if k.startswith('module'):
+                    name = k[7:]  # remove `module.`
+                else:
+                    name = k
+                new_state_dict[name] = v
+            model.load_state_dict(new_state_dict)
+        else:
+            model.load_state_dict(checkpoint)
+        print("=> Loaded checkpoint '{}'".format(checkpoint_path))
+    else:
+        print("=> Error: No checkpoint found at '{}'".format(checkpoint_path))
+        raise FileNotFoundError()
+
+
+def load_pretrained(model, url, filter_fn=None, strict=True):
+    if not url:
+        print("=> Warning: Pretrained model URL is empty, using random initialization.")
+        return
+
+    state_dict = load_state_dict_from_url(url, progress=False, map_location='cpu')
+
+    input_conv = 'conv_stem'
+    classifier = 'classifier'
+    in_chans = getattr(model, input_conv).weight.shape[1]
+    num_classes = getattr(model, classifier).weight.shape[0]
+
+    input_conv_weight = input_conv + '.weight'
+    pretrained_in_chans = state_dict[input_conv_weight].shape[1]
+    if in_chans != pretrained_in_chans:
+        if in_chans == 1:
+            print('=> Converting pretrained input conv {} from {} to 1 channel'.format(
+                input_conv_weight, pretrained_in_chans))
+            conv1_weight = state_dict[input_conv_weight]
+            state_dict[input_conv_weight] = conv1_weight.sum(dim=1, keepdim=True)
+        else:
+            print('=> Discarding pretrained input conv {} since input channel count != {}'.format(
+                input_conv_weight, pretrained_in_chans))
+            del state_dict[input_conv_weight]
+            strict = False
+
+    classifier_weight = classifier + '.weight'
+    pretrained_num_classes = state_dict[classifier_weight].shape[0]
+    if num_classes != pretrained_num_classes:
+        print('=> Discarding pretrained classifier since num_classes != {}'.format(pretrained_num_classes))
+        del state_dict[classifier_weight]
+        del state_dict[classifier + '.bias']
+        strict = False
+
+    if filter_fn is not None:
+        state_dict = filter_fn(state_dict)
+
+    model.load_state_dict(state_dict, strict=strict)

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/mobilenetv3.py

@@ -0,0 +1,366 @@
+""" MobileNet-V3
+
+A PyTorch impl of MobileNet-V3, compatible with TF weights from official impl.
+
+Paper: Searching for MobileNetV3 - https://arxiv.org/abs/1905.02244
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .activations import get_act_fn, get_act_layer, HardSwish
+from .config import layer_config_kwargs
+from .conv2d_layers import select_conv2d
+from .helpers import load_pretrained
+from .efficientnet_builder import (BN_EPS_TF_DEFAULT, EfficientNetBuilder, decode_arch_def,
+                                   initialize_weight_default, initialize_weight_goog,
+                                   resolve_act_layer, resolve_bn_args, round_channels)
+
+__all__ = ['mobilenetv3_rw', 'mobilenetv3_large_075', 'mobilenetv3_large_100', 'mobilenetv3_large_minimal_100',
+           'mobilenetv3_small_075', 'mobilenetv3_small_100', 'mobilenetv3_small_minimal_100',
+           'tf_mobilenetv3_large_075', 'tf_mobilenetv3_large_100', 'tf_mobilenetv3_large_minimal_100',
+           'tf_mobilenetv3_small_075', 'tf_mobilenetv3_small_100', 'tf_mobilenetv3_small_minimal_100']
+
+model_urls = {
+    'mobilenetv3_rw':
+        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mobilenetv3_100-35495452.pth',
+    'mobilenetv3_large_075': None,
+    'mobilenetv3_large_100':
+        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mobilenetv3_large_100_ra-f55367f5.pth',
+    'mobilenetv3_large_minimal_100': None,
+    'mobilenetv3_small_075': None,
+    'mobilenetv3_small_100': None,
+    'mobilenetv3_small_minimal_100': None,
+    'tf_mobilenetv3_large_075':
+        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_075-150ee8b0.pth',
+    'tf_mobilenetv3_large_100':
+        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_100-427764d5.pth',
+    'tf_mobilenetv3_large_minimal_100':
+        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_minimal_100-8596ae28.pth',
+    'tf_mobilenetv3_small_075':
+        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_075-da427f52.pth',
+    'tf_mobilenetv3_small_100':
+        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_100-37f49e2b.pth',
+    'tf_mobilenetv3_small_minimal_100':
+        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_minimal_100-922a7843.pth',
+}
+
+
+class MobileNetV3(nn.Module):
+    """ MobileNet-V3
+
+    A this model utilizes the MobileNet-v3 specific 'efficient head', where global pooling is done before the
+    head convolution without a final batch-norm layer before the classifier.
+
+    Paper: https://arxiv.org/abs/1905.02244
+    """
+
+    def __init__(self, block_args, num_classes=1000, in_chans=3, stem_size=16, num_features=1280, head_bias=True,
+                 channel_multiplier=1.0, pad_type='', act_layer=HardSwish, drop_rate=0., drop_connect_rate=0.,
+                 se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None, weight_init='goog'):
+        super(MobileNetV3, self).__init__()
+        self.drop_rate = drop_rate
+
+        stem_size = round_channels(stem_size, channel_multiplier)
+        self.conv_stem = select_conv2d(in_chans, stem_size, 3, stride=2, padding=pad_type)
+        self.bn1 = nn.BatchNorm2d(stem_size, **norm_kwargs)
+        self.act1 = act_layer(inplace=True)
+        in_chs = stem_size
+
+        builder = EfficientNetBuilder(
+            channel_multiplier, pad_type=pad_type, act_layer=act_layer, se_kwargs=se_kwargs,
+            norm_layer=norm_layer, norm_kwargs=norm_kwargs, drop_connect_rate=drop_connect_rate)
+        self.blocks = nn.Sequential(*builder(in_chs, block_args))
+        in_chs = builder.in_chs
+
+        self.global_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv_head = select_conv2d(in_chs, num_features, 1, padding=pad_type, bias=head_bias)
+        self.act2 = act_layer(inplace=True)
+        self.classifier = nn.Linear(num_features, num_classes)
+
+        for m in self.modules():
+            if weight_init == 'goog':
+                initialize_weight_goog(m)
+            else:
+                initialize_weight_default(m)
+
+    def as_sequential(self):
+        layers = [self.conv_stem, self.bn1, self.act1]
+        layers.extend(self.blocks)
+        layers.extend([
+            self.global_pool, self.conv_head, self.act2,
+            nn.Flatten(), nn.Dropout(self.drop_rate), self.classifier])
+        return nn.Sequential(*layers)
+
+    def features(self, x):
+        x = self.conv_stem(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+        x = self.blocks(x)
+        x = self.global_pool(x)
+        x = self.conv_head(x)
+        x = self.act2(x)
+        return x
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.flatten(1)
+        if self.drop_rate > 0.:
+            x = F.dropout(x, p=self.drop_rate, training=self.training)
+        return self.classifier(x)
+
+
+def _create_model(model_kwargs, variant, pretrained=False):
+    as_sequential = model_kwargs.pop('as_sequential', False)
+    model = MobileNetV3(**model_kwargs)
+    if pretrained and model_urls[variant]:
+        load_pretrained(model, model_urls[variant])
+    if as_sequential:
+        model = model.as_sequential()
+    return model
+
+
+def _gen_mobilenet_v3_rw(variant, channel_multiplier=1.0, pretrained=False, **kwargs):
+    """Creates a MobileNet-V3 model (RW variant).
+
+    Paper: https://arxiv.org/abs/1905.02244
+
+    This was my first attempt at reproducing the MobileNet-V3 from paper alone. It came close to the
+    eventual Tensorflow reference impl but has a few differences:
+    1. This model has no bias on the head convolution
+    2. This model forces no residual (noskip) on the first DWS block, this is different than MnasNet
+    3. This model always uses ReLU for the SE activation layer, other models in the family inherit their act layer
+       from their parent block
+    4. This model does not enforce divisible by 8 limitation on the SE reduction channel count
+
+    Overall the changes are fairly minor and result in a very small parameter count difference and no
+    top-1/5
+
+    Args:
+      channel_multiplier: multiplier to number of channels per layer.
+    """
+    arch_def = [
+        # stage 0, 112x112 in
+        ['ds_r1_k3_s1_e1_c16_nre_noskip'],  # relu
+        # stage 1, 112x112 in
+        ['ir_r1_k3_s2_e4_c24_nre', 'ir_r1_k3_s1_e3_c24_nre'],  # relu
+        # stage 2, 56x56 in
+        ['ir_r3_k5_s2_e3_c40_se0.25_nre'],  # relu
+        # stage 3, 28x28 in
+        ['ir_r1_k3_s2_e6_c80', 'ir_r1_k3_s1_e2.5_c80', 'ir_r2_k3_s1_e2.3_c80'],  # hard-swish
+        # stage 4, 14x14in
+        ['ir_r2_k3_s1_e6_c112_se0.25'],  # hard-swish
+        # stage 5, 14x14in
+        ['ir_r3_k5_s2_e6_c160_se0.25'],  # hard-swish
+        # stage 6, 7x7 in
+        ['cn_r1_k1_s1_c960'],  # hard-swish
+    ]
+    with layer_config_kwargs(kwargs):
+        model_kwargs = dict(
+            block_args=decode_arch_def(arch_def),
+            head_bias=False,  # one of my mistakes
+            channel_multiplier=channel_multiplier,
+            act_layer=resolve_act_layer(kwargs, 'hard_swish'),
+            se_kwargs=dict(gate_fn=get_act_fn('hard_sigmoid'), reduce_mid=True),
+            norm_kwargs=resolve_bn_args(kwargs),
+            **kwargs,
+        )
+        model = _create_model(model_kwargs, variant, pretrained)
+    return model
+
+
+def _gen_mobilenet_v3(variant, channel_multiplier=1.0, pretrained=False, **kwargs):
+    """Creates a MobileNet-V3 large/small/minimal models.
+
+    Ref impl: https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet_v3.py
+    Paper: https://arxiv.org/abs/1905.02244
+
+    Args:
+      channel_multiplier: multiplier to number of channels per layer.
+    """
+    if 'small' in variant:
+        num_features = 1024
+        if 'minimal' in variant:
+            act_layer = 'relu'
+            arch_def = [
+                # stage 0, 112x112 in
+                ['ds_r1_k3_s2_e1_c16'],
+                # stage 1, 56x56 in
+                ['ir_r1_k3_s2_e4.5_c24', 'ir_r1_k3_s1_e3.67_c24'],
+                # stage 2, 28x28 in
+                ['ir_r1_k3_s2_e4_c40', 'ir_r2_k3_s1_e6_c40'],
+                # stage 3, 14x14 in
+                ['ir_r2_k3_s1_e3_c48'],
+                # stage 4, 14x14in
+                ['ir_r3_k3_s2_e6_c96'],
+                # stage 6, 7x7 in
+                ['cn_r1_k1_s1_c576'],
+            ]
+        else:
+            act_layer = 'hard_swish'
+            arch_def = [
+                # stage 0, 112x112 in
+                ['ds_r1_k3_s2_e1_c16_se0.25_nre'],  # relu
+                # stage 1, 56x56 in
+                ['ir_r1_k3_s2_e4.5_c24_nre', 'ir_r1_k3_s1_e3.67_c24_nre'],  # relu
+                # stage 2, 28x28 in
+                ['ir_r1_k5_s2_e4_c40_se0.25', 'ir_r2_k5_s1_e6_c40_se0.25'],  # hard-swish
+                # stage 3, 14x14 in
+                ['ir_r2_k5_s1_e3_c48_se0.25'],  # hard-swish
+                # stage 4, 14x14in
+                ['ir_r3_k5_s2_e6_c96_se0.25'],  # hard-swish
+                # stage 6, 7x7 in
+                ['cn_r1_k1_s1_c576'],  # hard-swish
+            ]
+    else:
+        num_features = 1280
+        if 'minimal' in variant:
+            act_layer = 'relu'
+            arch_def = [
+                # stage 0, 112x112 in
+                ['ds_r1_k3_s1_e1_c16'],
+                # stage 1, 112x112 in
+                ['ir_r1_k3_s2_e4_c24', 'ir_r1_k3_s1_e3_c24'],
+                # stage 2, 56x56 in
+                ['ir_r3_k3_s2_e3_c40'],
+                # stage 3, 28x28 in
+                ['ir_r1_k3_s2_e6_c80', 'ir_r1_k3_s1_e2.5_c80', 'ir_r2_k3_s1_e2.3_c80'],
+                # stage 4, 14x14in
+                ['ir_r2_k3_s1_e6_c112'],
+                # stage 5, 14x14in
+                ['ir_r3_k3_s2_e6_c160'],
+                # stage 6, 7x7 in
+                ['cn_r1_k1_s1_c960'],
+            ]
+        else:
+            act_layer = 'hard_swish'
+            arch_def = [
+                # stage 0, 112x112 in
+                ['ds_r1_k3_s1_e1_c16_nre'],  # relu
+                # stage 1, 112x112 in
+                ['ir_r1_k3_s2_e4_c24_nre', 'ir_r1_k3_s1_e3_c24_nre'],  # relu
+                # stage 2, 56x56 in
+                ['ir_r3_k5_s2_e3_c40_se0.25_nre'],  # relu
+                # stage 3, 28x28 in
+                ['ir_r1_k3_s2_e6_c80', 'ir_r1_k3_s1_e2.5_c80', 'ir_r2_k3_s1_e2.3_c80'],  # hard-swish
+                # stage 4, 14x14in
+                ['ir_r2_k3_s1_e6_c112_se0.25'],  # hard-swish
+                # stage 5, 14x14in
+                ['ir_r3_k5_s2_e6_c160_se0.25'],  # hard-swish
+                # stage 6, 7x7 in
+                ['cn_r1_k1_s1_c960'],  # hard-swish
+            ]
+    with layer_config_kwargs(kwargs):
+        model_kwargs = dict(
+            block_args=decode_arch_def(arch_def),
+            num_features=num_features,
+            stem_size=16,
+            channel_multiplier=channel_multiplier,
+            act_layer=resolve_act_layer(kwargs, act_layer),
+            se_kwargs=dict(
+                act_layer=get_act_layer('relu'), gate_fn=get_act_fn('hard_sigmoid'), reduce_mid=True, divisor=8),
+            norm_kwargs=resolve_bn_args(kwargs),
+            **kwargs,
+        )
+        model = _create_model(model_kwargs, variant, pretrained)
+    return model
+
+
+def mobilenetv3_rw(pretrained=False, **kwargs):
+    """ MobileNet-V3 RW
+    Attn: See note in gen function for this variant.
+    """
+    # NOTE for train set drop_rate=0.2
+    if pretrained:
+        # pretrained model trained with non-default BN epsilon
+        kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
+    model = _gen_mobilenet_v3_rw('mobilenetv3_rw', 1.0, pretrained=pretrained, **kwargs)
+    return model
+
+
+def mobilenetv3_large_075(pretrained=False, **kwargs):
+    """ MobileNet V3 Large 0.75"""
+    # NOTE for train set drop_rate=0.2
+    model = _gen_mobilenet_v3('mobilenetv3_large_075', 0.75, pretrained=pretrained, **kwargs)
+    return model
+
+
+def mobilenetv3_large_100(pretrained=False, **kwargs):
+    """ MobileNet V3 Large 1.0 """
+    # NOTE for train set drop_rate=0.2
+    model = _gen_mobilenet_v3('mobilenetv3_large_100', 1.0, pretrained=pretrained, **kwargs)
+    return model
+
+
+def mobilenetv3_large_minimal_100(pretrained=False, **kwargs):
+    """ MobileNet V3 Large (Minimalistic) 1.0 """
+    # NOTE for train set drop_rate=0.2
+    model = _gen_mobilenet_v3('mobilenetv3_large_minimal_100', 1.0, pretrained=pretrained, **kwargs)
+    return model
+
+
+def mobilenetv3_small_075(pretrained=False, **kwargs):
+    """ MobileNet V3 Small 0.75 """
+    model = _gen_mobilenet_v3('mobilenetv3_small_075', 0.75, pretrained=pretrained, **kwargs)
+    return model
+
+
+def mobilenetv3_small_100(pretrained=False, **kwargs):
+    """ MobileNet V3 Small 1.0 """
+    model = _gen_mobilenet_v3('mobilenetv3_small_100', 1.0, pretrained=pretrained, **kwargs)
+    return model
+
+
+def mobilenetv3_small_minimal_100(pretrained=False, **kwargs):
+    """ MobileNet V3 Small (Minimalistic) 1.0 """
+    model = _gen_mobilenet_v3('mobilenetv3_small_minimal_100', 1.0, pretrained=pretrained, **kwargs)
+    return model
+
+
+def tf_mobilenetv3_large_075(pretrained=False, **kwargs):
+    """ MobileNet V3 Large 0.75. Tensorflow compat variant. """
+    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
+    kwargs['pad_type'] = 'same'
+    model = _gen_mobilenet_v3('tf_mobilenetv3_large_075', 0.75, pretrained=pretrained, **kwargs)
+    return model
+
+
+def tf_mobilenetv3_large_100(pretrained=False, **kwargs):
+    """ MobileNet V3 Large 1.0. Tensorflow compat variant. """
+    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
+    kwargs['pad_type'] = 'same'
+    model = _gen_mobilenet_v3('tf_mobilenetv3_large_100', 1.0, pretrained=pretrained, **kwargs)
+    return model
+
+
+def tf_mobilenetv3_large_minimal_100(pretrained=False, **kwargs):
+    """ MobileNet V3 Large Minimalistic 1.0. Tensorflow compat variant. """
+    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
+    kwargs['pad_type'] = 'same'
+    model = _gen_mobilenet_v3('tf_mobilenetv3_large_minimal_100', 1.0, pretrained=pretrained, **kwargs)
+    return model
+
+
+def tf_mobilenetv3_small_075(pretrained=False, **kwargs):
+    """ MobileNet V3 Small 0.75. Tensorflow compat variant. """
+    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
+    kwargs['pad_type'] = 'same'
+    model = _gen_mobilenet_v3('tf_mobilenetv3_small_075', 0.75, pretrained=pretrained, **kwargs)
+    return model
+
+
+def tf_mobilenetv3_small_100(pretrained=False, **kwargs):
+    """ MobileNet V3 Small 1.0. Tensorflow compat variant."""
+    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
+    kwargs['pad_type'] = 'same'
+    model = _gen_mobilenet_v3('tf_mobilenetv3_small_100', 1.0, pretrained=pretrained, **kwargs)
+    return model
+
+
+def tf_mobilenetv3_small_minimal_100(pretrained=False, **kwargs):
+    """ MobileNet V3 Small Minimalistic 1.0. Tensorflow compat variant. """
+    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
+    kwargs['pad_type'] = 'same'
+    model = _gen_mobilenet_v3('tf_mobilenetv3_small_minimal_100', 1.0, pretrained=pretrained, **kwargs)
+    return model

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/model_factory.py

@@ -0,0 +1,27 @@
+from .config import set_layer_config
+from .helpers import load_checkpoint
+
+from .gen_efficientnet import *
+from .mobilenetv3 import *
+
+
+def create_model(
+        model_name='mnasnet_100',
+        pretrained=None,
+        num_classes=1000,
+        in_chans=3,
+        checkpoint_path='',
+        **kwargs):
+
+    model_kwargs = dict(num_classes=num_classes, in_chans=in_chans, pretrained=pretrained, **kwargs)
+
+    if model_name in globals():
+        create_fn = globals()[model_name]
+        model = create_fn(**model_kwargs)
+    else:
+        raise RuntimeError('Unknown model (%s)' % model_name)
+
+    if checkpoint_path and not pretrained:
+        load_checkpoint(model, checkpoint_path)
+
+    return model

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/version.py

@@ -0,0 +1 @@
+__version__ = '1.0.2'

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/hubconf.py

@@ -0,0 +1,84 @@
+dependencies = ['torch', 'math']
+
+from geffnet import efficientnet_b0
+from geffnet import efficientnet_b1
+from geffnet import efficientnet_b2
+from geffnet import efficientnet_b3
+
+from geffnet import efficientnet_es
+
+from geffnet import efficientnet_lite0
+
+from geffnet import mixnet_s
+from geffnet import mixnet_m
+from geffnet import mixnet_l
+from geffnet import mixnet_xl
+
+from geffnet import mobilenetv2_100
+from geffnet import mobilenetv2_110d
+from geffnet import mobilenetv2_120d
+from geffnet import mobilenetv2_140
+
+from geffnet import mobilenetv3_large_100
+from geffnet import mobilenetv3_rw
+from geffnet import mnasnet_a1
+from geffnet import mnasnet_b1
+from geffnet import fbnetc_100
+from geffnet import spnasnet_100
+
+from geffnet import tf_efficientnet_b0
+from geffnet import tf_efficientnet_b1
+from geffnet import tf_efficientnet_b2
+from geffnet import tf_efficientnet_b3
+from geffnet import tf_efficientnet_b4
+from geffnet import tf_efficientnet_b5
+from geffnet import tf_efficientnet_b6
+from geffnet import tf_efficientnet_b7
+from geffnet import tf_efficientnet_b8
+
+from geffnet import tf_efficientnet_b0_ap
+from geffnet import tf_efficientnet_b1_ap
+from geffnet import tf_efficientnet_b2_ap
+from geffnet import tf_efficientnet_b3_ap
+from geffnet import tf_efficientnet_b4_ap
+from geffnet import tf_efficientnet_b5_ap
+from geffnet import tf_efficientnet_b6_ap
+from geffnet import tf_efficientnet_b7_ap
+from geffnet import tf_efficientnet_b8_ap
+
+from geffnet import tf_efficientnet_b0_ns
+from geffnet import tf_efficientnet_b1_ns
+from geffnet import tf_efficientnet_b2_ns
+from geffnet import tf_efficientnet_b3_ns
+from geffnet import tf_efficientnet_b4_ns
+from geffnet import tf_efficientnet_b5_ns
+from geffnet import tf_efficientnet_b6_ns
+from geffnet import tf_efficientnet_b7_ns
+from geffnet import tf_efficientnet_l2_ns_475
+from geffnet import tf_efficientnet_l2_ns
+
+from geffnet import tf_efficientnet_es
+from geffnet import tf_efficientnet_em
+from geffnet import tf_efficientnet_el
+
+from geffnet import tf_efficientnet_cc_b0_4e
+from geffnet import tf_efficientnet_cc_b0_8e
+from geffnet import tf_efficientnet_cc_b1_8e
+
+from geffnet import tf_efficientnet_lite0
+from geffnet import tf_efficientnet_lite1
+from geffnet import tf_efficientnet_lite2
+from geffnet import tf_efficientnet_lite3
+from geffnet import tf_efficientnet_lite4
+
+from geffnet import tf_mixnet_s
+from geffnet import tf_mixnet_m
+from geffnet import tf_mixnet_l
+
+from geffnet import tf_mobilenetv3_large_075
+from geffnet import tf_mobilenetv3_large_100
+from geffnet import tf_mobilenetv3_large_minimal_100
+from geffnet import tf_mobilenetv3_small_075
+from geffnet import tf_mobilenetv3_small_100
+from geffnet import tf_mobilenetv3_small_minimal_100
+

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/onnx_export.py

@@ -0,0 +1,120 @@
+""" ONNX export script
+
+Export PyTorch models as ONNX graphs.
+
+This export script originally started as an adaptation of code snippets found at
+https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html
+
+The default parameters work with PyTorch 1.6 and ONNX 1.7 and produce an optimal ONNX graph
+for hosting in the ONNX runtime (see onnx_validate.py). To export an ONNX model compatible
+with caffe2 (see caffe2_benchmark.py and caffe2_validate.py), the --keep-init and --aten-fallback
+flags are currently required.
+
+Older versions of PyTorch/ONNX (tested PyTorch 1.4, ONNX 1.5) do not need extra flags for
+caffe2 compatibility, but they produce a model that isn't as fast running on ONNX runtime.
+
+Most new release of PyTorch and ONNX cause some sort of breakage in the export / usage of ONNX models.
+Please do your research and search ONNX and PyTorch issue tracker before asking me. Thanks.
+
+Copyright 2020 Ross Wightman
+"""
+import argparse
+import torch
+import numpy as np
+
+import onnx
+import geffnet
+
+parser = argparse.ArgumentParser(description='PyTorch ImageNet Validation')
+parser.add_argument('output', metavar='ONNX_FILE',
+                    help='output model filename')
+parser.add_argument('--model', '-m', metavar='MODEL', default='mobilenetv3_large_100',
+                    help='model architecture (default: mobilenetv3_large_100)')
+parser.add_argument('--opset', type=int, default=10,
+                    help='ONNX opset to use (default: 10)')
+parser.add_argument('--keep-init', action='store_true', default=False,
+                    help='Keep initializers as input. Needed for Caffe2 compatible export in newer PyTorch/ONNX.')
+parser.add_argument('--aten-fallback', action='store_true', default=False,
+                    help='Fallback to ATEN ops. Helps fix AdaptiveAvgPool issue with Caffe2 in newer PyTorch/ONNX.')
+parser.add_argument('--dynamic-size', action='store_true', default=False,
+                    help='Export model width dynamic width/height. Not recommended for "tf" models with SAME padding.')
+parser.add_argument('-b', '--batch-size', default=1, type=int,
+                    metavar='N', help='mini-batch size (default: 1)')
+parser.add_argument('--img-size', default=None, type=int,
+                    metavar='N', help='Input image dimension, uses model default if empty')
+parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN',
+                    help='Override mean pixel value of dataset')
+parser.add_argument('--std', type=float,  nargs='+', default=None, metavar='STD',
+                    help='Override std deviation of of dataset')
+parser.add_argument('--num-classes', type=int, default=1000,
+                    help='Number classes in dataset')
+parser.add_argument('--checkpoint', default='', type=str, metavar='PATH',
+                    help='path to checkpoint (default: none)')
+
+
+def main():
+    args = parser.parse_args()
+
+    args.pretrained = True
+    if args.checkpoint:
+        args.pretrained = False
+
+    print("==> Creating PyTorch {} model".format(args.model))
+    # NOTE exportable=True flag disables autofn/jit scripted activations and uses Conv2dSameExport layers
+    # for models using SAME padding
+    model = geffnet.create_model(
+        args.model,
+        num_classes=args.num_classes,
+        in_chans=3,
+        pretrained=args.pretrained,
+        checkpoint_path=args.checkpoint,
+        exportable=True)
+
+    model.eval()
+
+    example_input = torch.randn((args.batch_size, 3, args.img_size or 224, args.img_size or 224), requires_grad=True)
+
+    # Run model once before export trace, sets padding for models with Conv2dSameExport. This means
+    # that the padding for models with Conv2dSameExport (most models with tf_ prefix) is fixed for
+    # the input img_size specified in this script.
+    # Opset >= 11 should allow for dynamic padding, however I cannot get it to work due to
+    # issues in the tracing of the dynamic padding or errors attempting to export the model after jit
+    # scripting it (an approach that should work). Perhaps in a future PyTorch or ONNX versions...
+    model(example_input)
+
+    print("==> Exporting model to ONNX format at '{}'".format(args.output))
+    input_names = ["input0"]
+    output_names = ["output0"]
+    dynamic_axes = {'input0': {0: 'batch'}, 'output0': {0: 'batch'}}
+    if args.dynamic_size:
+        dynamic_axes['input0'][2] = 'height'
+        dynamic_axes['input0'][3] = 'width'
+    if args.aten_fallback:
+        export_type = torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK
+    else:
+        export_type = torch.onnx.OperatorExportTypes.ONNX
+
+    torch_out = torch.onnx._export(
+        model, example_input, args.output, export_params=True, verbose=True, input_names=input_names,
+        output_names=output_names, keep_initializers_as_inputs=args.keep_init, dynamic_axes=dynamic_axes,
+        opset_version=args.opset, operator_export_type=export_type)
+
+    print("==> Loading and checking exported model from '{}'".format(args.output))
+    onnx_model = onnx.load(args.output)
+    onnx.checker.check_model(onnx_model)  # assuming throw on error
+    print("==> Passed")
+
+    if args.keep_init and args.aten_fallback:
+        import caffe2.python.onnx.backend as onnx_caffe2
+        # Caffe2 loading only works properly in newer PyTorch/ONNX combos when
+        # keep_initializers_as_inputs and aten_fallback are set to True.
+        print("==> Loading model into Caffe2 backend and comparing forward pass.".format(args.output))
+        caffe2_backend = onnx_caffe2.prepare(onnx_model)
+        B = {onnx_model.graph.input[0].name: x.data.numpy()}
+        c2_out = caffe2_backend.run(B)[0]
+        np.testing.assert_almost_equal(torch_out.data.numpy(), c2_out, decimal=5)
+        print("==> Passed")
+
+
+if __name__ == '__main__':
+    main()

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/onnx_optimize.py

@@ -0,0 +1,84 @@
+""" ONNX optimization script
+
+Run ONNX models through the optimizer to prune unneeded nodes, fuse batchnorm layers into conv, etc.
+
+NOTE: This isn't working consistently in recent PyTorch/ONNX combos (ie PyTorch 1.6 and ONNX 1.7),
+it seems time to switch to using the onnxruntime online optimizer (can also be saved for offline).
+
+Copyright 2020 Ross Wightman
+"""
+import argparse
+import warnings
+
+import onnx
+from onnx import optimizer
+
+
+parser = argparse.ArgumentParser(description="Optimize ONNX model")
+
+parser.add_argument("model", help="The ONNX model")
+parser.add_argument("--output", required=True, help="The optimized model output filename")
+
+
+def traverse_graph(graph, prefix=''):
+    content = []
+    indent = prefix + '  '
+    graphs = []
+    num_nodes = 0
+    for node in graph.node:
+        pn, gs = onnx.helper.printable_node(node, indent, subgraphs=True)
+        assert isinstance(gs, list)
+        content.append(pn)
+        graphs.extend(gs)
+        num_nodes += 1
+    for g in graphs:
+        g_count, g_str = traverse_graph(g)
+        content.append('\n' + g_str)
+        num_nodes += g_count
+    return num_nodes, '\n'.join(content)
+
+
+def main():
+    args = parser.parse_args()
+    onnx_model = onnx.load(args.model)
+    num_original_nodes, original_graph_str = traverse_graph(onnx_model.graph)
+
+    # Optimizer passes to perform
+    passes = [
+        #'eliminate_deadend',
+        'eliminate_identity',
+        'eliminate_nop_dropout',
+        'eliminate_nop_pad',
+        'eliminate_nop_transpose',
+        'eliminate_unused_initializer',
+        'extract_constant_to_initializer',
+        'fuse_add_bias_into_conv',
+        'fuse_bn_into_conv',
+        'fuse_consecutive_concats',
+        'fuse_consecutive_reduce_unsqueeze',
+        'fuse_consecutive_squeezes',
+        'fuse_consecutive_transposes',
+        #'fuse_matmul_add_bias_into_gemm',
+        'fuse_pad_into_conv',
+        #'fuse_transpose_into_gemm',
+        #'lift_lexical_references',
+    ]
+
+    # Apply the optimization on the original serialized model
+    # WARNING I've had issues with optimizer in recent versions of PyTorch / ONNX causing
+    # 'duplicate definition of name' errors, see: https://github.com/onnx/onnx/issues/2401
+    # It may be better to rely on onnxruntime optimizations, see onnx_validate.py script.
+    warnings.warn("I've had issues with optimizer in recent versions of PyTorch / ONNX."
+                  "Try onnxruntime optimization if this doesn't work.")
+    optimized_model = optimizer.optimize(onnx_model, passes)
+
+    num_optimized_nodes, optimzied_graph_str = traverse_graph(optimized_model.graph)
+    print('==> The model after optimization:\n{}\n'.format(optimzied_graph_str))
+    print('==> The optimized model has {} nodes, the original had {}.'.format(num_optimized_nodes, num_original_nodes))
+
+    # Save the ONNX model
+    onnx.save(optimized_model, args.output)
+
+
+if __name__ == "__main__":
+    main()

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/onnx_to_caffe.py

@@ -0,0 +1,27 @@
+import argparse
+
+import onnx
+from caffe2.python.onnx.backend import Caffe2Backend
+
+
+parser = argparse.ArgumentParser(description="Convert ONNX to Caffe2")
+
+parser.add_argument("model", help="The ONNX model")
+parser.add_argument("--c2-prefix", required=True,
+    help="The output file prefix for the caffe2 model init and predict file. ")
+
+
+def main():
+    args = parser.parse_args()
+    onnx_model = onnx.load(args.model)
+    caffe2_init, caffe2_predict = Caffe2Backend.onnx_graph_to_caffe2_net(onnx_model)
+    caffe2_init_str = caffe2_init.SerializeToString()
+    with open(args.c2_prefix + '.init.pb', "wb") as f:
+        f.write(caffe2_init_str)
+    caffe2_predict_str = caffe2_predict.SerializeToString()
+    with open(args.c2_prefix + '.predict.pb', "wb") as f:
+        f.write(caffe2_predict_str)
+
+
+if __name__ == "__main__":
+    main()

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/onnx_validate.py

@@ -0,0 +1,112 @@
+""" ONNX-runtime validation script
+
+This script was created to verify accuracy and performance of exported ONNX
+models running with the onnxruntime. It utilizes the PyTorch dataloader/processing
+pipeline for a fair comparison against the originals.
+
+Copyright 2020 Ross Wightman
+"""
+import argparse
+import numpy as np
+import onnxruntime
+from data import create_loader, resolve_data_config, Dataset
+from utils import AverageMeter
+import time
+
+parser = argparse.ArgumentParser(description='Caffe2 ImageNet Validation')
+parser.add_argument('data', metavar='DIR',
+                    help='path to dataset')
+parser.add_argument('--onnx-input', default='', type=str, metavar='PATH',
+                    help='path to onnx model/weights file')
+parser.add_argument('--onnx-output-opt', default='', type=str, metavar='PATH',
+                    help='path to output optimized onnx graph')
+parser.add_argument('--profile', action='store_true', default=False,
+                    help='Enable profiler output.')
+parser.add_argument('-j', '--workers', default=2, type=int, metavar='N',
+                    help='number of data loading workers (default: 2)')
+parser.add_argument('-b', '--batch-size', default=256, type=int,
+                    metavar='N', help='mini-batch size (default: 256)')
+parser.add_argument('--img-size', default=None, type=int,
+                    metavar='N', help='Input image dimension, uses model default if empty')
+parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN',
+                    help='Override mean pixel value of dataset')
+parser.add_argument('--std', type=float,  nargs='+', default=None, metavar='STD',
+                    help='Override std deviation of of dataset')
+parser.add_argument('--crop-pct', type=float, default=None, metavar='PCT',
+                    help='Override default crop pct of 0.875')
+parser.add_argument('--interpolation', default='', type=str, metavar='NAME',
+                    help='Image resize interpolation type (overrides model)')
+parser.add_argument('--tf-preprocessing', dest='tf_preprocessing', action='store_true',
+                    help='use tensorflow mnasnet preporcessing')
+parser.add_argument('--print-freq', '-p', default=10, type=int,
+                    metavar='N', help='print frequency (default: 10)')
+
+
+def main():
+    args = parser.parse_args()
+    args.gpu_id = 0
+
+    # Set graph optimization level
+    sess_options = onnxruntime.SessionOptions()
+    sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+    if args.profile:
+        sess_options.enable_profiling = True
+    if args.onnx_output_opt:
+        sess_options.optimized_model_filepath = args.onnx_output_opt
+
+    session = onnxruntime.InferenceSession(args.onnx_input, sess_options)
+
+    data_config = resolve_data_config(None, args)
+    loader = create_loader(
+        Dataset(args.data, load_bytes=args.tf_preprocessing),
+        input_size=data_config['input_size'],
+        batch_size=args.batch_size,
+        use_prefetcher=False,
+        interpolation=data_config['interpolation'],
+        mean=data_config['mean'],
+        std=data_config['std'],
+        num_workers=args.workers,
+        crop_pct=data_config['crop_pct'],
+        tensorflow_preprocessing=args.tf_preprocessing)
+
+    input_name = session.get_inputs()[0].name
+
+    batch_time = AverageMeter()
+    top1 = AverageMeter()
+    top5 = AverageMeter()
+    end = time.time()
+    for i, (input, target) in enumerate(loader):
+        # run the net and return prediction
+        output = session.run([], {input_name: input.data.numpy()})
+        output = output[0]
+
+        # measure accuracy and record loss
+        prec1, prec5 = accuracy_np(output, target.numpy())
+        top1.update(prec1.item(), input.size(0))
+        top5.update(prec5.item(), input.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if i % args.print_freq == 0:
+            print('Test: [{0}/{1}]\t'
+                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f}, {rate_avg:.3f}/s, {ms_avg:.3f} ms/sample) \t'
+                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
+                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
+                i, len(loader), batch_time=batch_time, rate_avg=input.size(0) / batch_time.avg,
+                ms_avg=100 * batch_time.avg / input.size(0), top1=top1, top5=top5))
+
+    print(' * Prec@1 {top1.avg:.3f} ({top1a:.3f}) Prec@5 {top5.avg:.3f} ({top5a:.3f})'.format(
+        top1=top1, top1a=100-top1.avg, top5=top5, top5a=100.-top5.avg))
+
+
+def accuracy_np(output, target):
+    max_indices = np.argsort(output, axis=1)[:, ::-1]
+    top5 = 100 * np.equal(max_indices[:, :5], target[:, np.newaxis]).sum(axis=1).mean()
+    top1 = 100 * np.equal(max_indices[:, 0], target).mean()
+    return top1, top5
+
+
+if __name__ == '__main__':
+    main()

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/requirements.txt

@@ -0,0 +1,2 @@
+torch>=1.2.0
+torchvision>=0.4.0

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/setup.py

@@ -0,0 +1,47 @@
+""" Setup
+"""
+from setuptools import setup, find_packages
+from codecs import open
+from os import path
+
+here = path.abspath(path.dirname(__file__))
+
+# Get the long description from the README file
+with open(path.join(here, 'README.md'), encoding='utf-8') as f:
+    long_description = f.read()
+
+exec(open('geffnet/version.py').read())
+setup(
+    name='geffnet',
+    version=__version__,
+    description='(Generic) EfficientNets for PyTorch',
+    long_description=long_description,
+    long_description_content_type='text/markdown',
+    url='https://github.com/rwightman/gen-efficientnet-pytorch',
+    author='Ross Wightman',
+    author_email='hello@rwightman.com',
+    classifiers=[
+        # How mature is this project? Common values are
+        #   3 - Alpha
+        #   4 - Beta
+        #   5 - Production/Stable
+        'Development Status :: 3 - Alpha',
+        'Intended Audience :: Education',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: Apache Software License',
+        'Programming Language :: Python :: 3.6',
+        'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
+        'Topic :: Scientific/Engineering',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Software Development',
+        'Topic :: Software Development :: Libraries',
+        'Topic :: Software Development :: Libraries :: Python Modules',
+    ],
+
+    # Note that this is a string of words separated by whitespace, not a list.
+    keywords='pytorch pretrained models efficientnet mixnet mobilenetv3 mnasnet',
+    packages=find_packages(exclude=['data']),
+    install_requires=['torch >= 1.4', 'torchvision'],
+    python_requires='>=3.6',
+)

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/utils.py

@@ -0,0 +1,52 @@
+import os
+
+
+class AverageMeter:
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the precision@k for the specified values of k"""
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].reshape(-1).float().sum(0)
+        res.append(correct_k.mul_(100.0 / batch_size))
+    return res
+
+
+def get_outdir(path, *paths, inc=False):
+    outdir = os.path.join(path, *paths)
+    if not os.path.exists(outdir):
+        os.makedirs(outdir)
+    elif inc:
+        count = 1
+        outdir_inc = outdir + '-' + str(count)
+        while os.path.exists(outdir_inc):
+            count = count + 1
+            outdir_inc = outdir + '-' + str(count)
+            assert count < 100
+        outdir = outdir_inc
+        os.makedirs(outdir)
+    return outdir
+

invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/validate.py

@@ -0,0 +1,166 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import time
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+from contextlib import suppress
+
+import geffnet
+from data import Dataset, create_loader, resolve_data_config
+from utils import accuracy, AverageMeter
+
+has_native_amp = False
+try:
+    if getattr(torch.cuda.amp, 'autocast') is not None:
+        has_native_amp = True
+except AttributeError:
+    pass
+
+torch.backends.cudnn.benchmark = True
+
+parser = argparse.ArgumentParser(description='PyTorch ImageNet Validation')
+parser.add_argument('data', metavar='DIR',
+                    help='path to dataset')
+parser.add_argument('--model', '-m', metavar='MODEL', default='spnasnet1_00',
+                    help='model architecture (default: dpn92)')
+parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
+                    help='number of data loading workers (default: 2)')
+parser.add_argument('-b', '--batch-size', default=256, type=int,
+                    metavar='N', help='mini-batch size (default: 256)')
+parser.add_argument('--img-size', default=None, type=int,
+                    metavar='N', help='Input image dimension, uses model default if empty')
+parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN',
+                    help='Override mean pixel value of dataset')
+parser.add_argument('--std', type=float,  nargs='+', default=None, metavar='STD',
+                    help='Override std deviation of of dataset')
+parser.add_argument('--crop-pct', type=float, default=None, metavar='PCT',
+                    help='Override default crop pct of 0.875')
+parser.add_argument('--interpolation', default='', type=str, metavar='NAME',
+                    help='Image resize interpolation type (overrides model)')
+parser.add_argument('--num-classes', type=int, default=1000,
+                    help='Number classes in dataset')
+parser.add_argument('--print-freq', '-p', default=10, type=int,
+                    metavar='N', help='print frequency (default: 10)')
+parser.add_argument('--checkpoint', default='', type=str, metavar='PATH',
+                    help='path to latest checkpoint (default: none)')
+parser.add_argument('--pretrained', dest='pretrained', action='store_true',
+                    help='use pre-trained model')
+parser.add_argument('--torchscript', dest='torchscript', action='store_true',
+                    help='convert model torchscript for inference')
+parser.add_argument('--num-gpu', type=int, default=1,
+                    help='Number of GPUS to use')
+parser.add_argument('--tf-preprocessing', dest='tf_preprocessing', action='store_true',
+                    help='use tensorflow mnasnet preporcessing')
+parser.add_argument('--no-cuda', dest='no_cuda', action='store_true',
+                    help='')
+parser.add_argument('--channels-last', action='store_true', default=False,
+                    help='Use channels_last memory layout')
+parser.add_argument('--amp', action='store_true', default=False,
+                    help='Use native Torch AMP mixed precision.')
+
+
+def main():
+    args = parser.parse_args()
+
+    if not args.checkpoint and not args.pretrained:
+        args.pretrained = True
+
+    amp_autocast = suppress  # do nothing
+    if args.amp:
+        if not has_native_amp:
+            print("Native Torch AMP is not available (requires torch >= 1.6), using FP32.")
+        else:
+            amp_autocast = torch.cuda.amp.autocast
+
+    # create model
+    model = geffnet.create_model(
+        args.model,
+        num_classes=args.num_classes,
+        in_chans=3,
+        pretrained=args.pretrained,
+        checkpoint_path=args.checkpoint,
+        scriptable=args.torchscript)
+
+    if args.channels_last:
+        model = model.to(memory_format=torch.channels_last)
+
+    if args.torchscript:
+        torch.jit.optimized_execution(True)
+        model = torch.jit.script(model)
+
+    print('Model %s created, param count: %d' %
+          (args.model, sum([m.numel() for m in model.parameters()])))
+
+    data_config = resolve_data_config(model, args)
+
+    criterion = nn.CrossEntropyLoss()
+
+    if not args.no_cuda:
+        if args.num_gpu > 1:
+            model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu))).cuda()
+        else:
+            model = model.cuda()
+        criterion = criterion.cuda()
+
+    loader = create_loader(
+        Dataset(args.data, load_bytes=args.tf_preprocessing),
+        input_size=data_config['input_size'],
+        batch_size=args.batch_size,
+        use_prefetcher=not args.no_cuda,
+        interpolation=data_config['interpolation'],
+        mean=data_config['mean'],
+        std=data_config['std'],
+        num_workers=args.workers,
+        crop_pct=data_config['crop_pct'],
+        tensorflow_preprocessing=args.tf_preprocessing)
+
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    top1 = AverageMeter()
+    top5 = AverageMeter()
+
+    model.eval()
+    end = time.time()
+    with torch.no_grad():
+        for i, (input, target) in enumerate(loader):
+            if not args.no_cuda:
+                target = target.cuda()
+                input = input.cuda()
+            if args.channels_last:
+                input = input.contiguous(memory_format=torch.channels_last)
+
+            # compute output
+            with amp_autocast():
+                output = model(input)
+                loss = criterion(output, target)
+
+            # measure accuracy and record loss
+            prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
+            losses.update(loss.item(), input.size(0))
+            top1.update(prec1.item(), input.size(0))
+            top5.update(prec5.item(), input.size(0))
+
+            # measure elapsed time
+            batch_time.update(time.time() - end)
+            end = time.time()
+
+            if i % args.print_freq == 0:
+                print('Test: [{0}/{1}]\t'
+                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f}, {rate_avg:.3f}/s) \t'
+                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
+                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
+                      'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
+                    i, len(loader), batch_time=batch_time,
+                    rate_avg=input.size(0) / batch_time.avg,
+                    loss=losses, top1=top1, top5=top5))
+
+    print(' * Prec@1 {top1.avg:.3f} ({top1a:.3f}) Prec@5 {top5.avg:.3f} ({top5a:.3f})'.format(
+        top1=top1, top1a=100-top1.avg, top5=top5, top5a=100.-top5.avg))
+
+
+if __name__ == '__main__':
+    main()

invokeai/backend/image_util/normal_bae/nets/submodules/encoder.py

@@ -0,0 +1,34 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Encoder(nn.Module):
+    def __init__(self):
+        super(Encoder, self).__init__()
+
+        basemodel_name = 'tf_efficientnet_b5_ap'
+        print('Loading base model ()...'.format(basemodel_name), end='')
+        repo_path = os.path.join(os.path.dirname(__file__), 'efficientnet_repo')
+        basemodel = torch.hub.load(repo_path, basemodel_name, pretrained=False, source='local')
+        print('Done.')
+
+        # Remove last layer
+        print('Removing last two layers (global_pool & classifier).')
+        basemodel.global_pool = nn.Identity()
+        basemodel.classifier = nn.Identity()
+
+        self.original_model = basemodel
+
+    def forward(self, x):
+        features = [x]
+        for k, v in self.original_model._modules.items():
+            if (k == 'blocks'):
+                for ki, vi in v._modules.items():
+                    features.append(vi(features[-1]))
+            else:
+                features.append(v(features[-1]))
+        return features
+
+

invokeai/backend/image_util/normal_bae/nets/submodules/submodules.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+########################################################################################################################
+
+
+# Upsample + BatchNorm
+class UpSampleBN(nn.Module):
+    def __init__(self, skip_input, output_features):
+        super(UpSampleBN, self).__init__()
+
+        self._net = nn.Sequential(nn.Conv2d(skip_input, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.BatchNorm2d(output_features),
+                                  nn.LeakyReLU(),
+                                  nn.Conv2d(output_features, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.BatchNorm2d(output_features),
+                                  nn.LeakyReLU())
+
+    def forward(self, x, concat_with):
+        up_x = F.interpolate(x, size=[concat_with.size(2), concat_with.size(3)], mode='bilinear', align_corners=True)
+        f = torch.cat([up_x, concat_with], dim=1)
+        return self._net(f)
+
+
+# Upsample + GroupNorm + Weight Standardization
+class UpSampleGN(nn.Module):
+    def __init__(self, skip_input, output_features):
+        super(UpSampleGN, self).__init__()
+
+        self._net = nn.Sequential(Conv2d(skip_input, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.GroupNorm(8, output_features),
+                                  nn.LeakyReLU(),
+                                  Conv2d(output_features, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.GroupNorm(8, output_features),
+                                  nn.LeakyReLU())
+
+    def forward(self, x, concat_with):
+        up_x = F.interpolate(x, size=[concat_with.size(2), concat_with.size(3)], mode='bilinear', align_corners=True)
+        f = torch.cat([up_x, concat_with], dim=1)
+        return self._net(f)
+
+
+# Conv2d with weight standardization
+class Conv2d(nn.Conv2d):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                 padding=0, dilation=1, groups=1, bias=True):
+        super(Conv2d, self).__init__(in_channels, out_channels, kernel_size, stride,
+                 padding, dilation, groups, bias)
+
+    def forward(self, x):
+        weight = self.weight
+        weight_mean = weight.mean(dim=1, keepdim=True).mean(dim=2,
+                                  keepdim=True).mean(dim=3, keepdim=True)
+        weight = weight - weight_mean
+        std = weight.view(weight.size(0), -1).std(dim=1).view(-1, 1, 1, 1) + 1e-5
+        weight = weight / std.expand_as(weight)
+        return F.conv2d(x, weight, self.bias, self.stride,
+                        self.padding, self.dilation, self.groups)
+
+
+# normalize
+def norm_normalize(norm_out):
+    min_kappa = 0.01
+    norm_x, norm_y, norm_z, kappa = torch.split(norm_out, 1, dim=1)
+    norm = torch.sqrt(norm_x ** 2.0 + norm_y ** 2.0 + norm_z ** 2.0) + 1e-10
+    kappa = F.elu(kappa) + 1.0 + min_kappa
+    final_out = torch.cat([norm_x / norm, norm_y / norm, norm_z / norm, kappa], dim=1)
+    return final_out
+
+
+# uncertainty-guided sampling (only used during training)
+@torch.no_grad()
+def sample_points(init_normal, gt_norm_mask, sampling_ratio, beta):
+    device = init_normal.device
+    B, _, H, W = init_normal.shape
+    N = int(sampling_ratio * H * W)
+    beta = beta
+
+    # uncertainty map
+    uncertainty_map = -1 * init_normal[:, 3, :, :]  # B, H, W
+
+    # gt_invalid_mask (B, H, W)
+    if gt_norm_mask is not None:
+        gt_invalid_mask = F.interpolate(gt_norm_mask.float(), size=[H, W], mode='nearest')
+        gt_invalid_mask = gt_invalid_mask[:, 0, :, :] < 0.5
+        uncertainty_map[gt_invalid_mask] = -1e4
+
+    # (B, H*W)
+    _, idx = uncertainty_map.view(B, -1).sort(1, descending=True)
+
+    # importance sampling
+    if int(beta * N) > 0:
+        importance = idx[:, :int(beta * N)]    # B, beta*N
+
+        # remaining
+        remaining = idx[:, int(beta * N):]     # B, H*W - beta*N
+
+        # coverage
+        num_coverage = N - int(beta * N)
+
+        if num_coverage <= 0:
+            samples = importance
+        else:
+            coverage_list = []
+            for i in range(B):
+                idx_c = torch.randperm(remaining.size()[1])    # shuffles "H*W - beta*N"
+                coverage_list.append(remaining[i, :][idx_c[:num_coverage]].view(1, -1))     # 1, N-beta*N
+            coverage = torch.cat(coverage_list, dim=0)                                      # B, N-beta*N
+            samples = torch.cat((importance, coverage), dim=1)                              # B, N
+
+    else:
+        # remaining
+        remaining = idx[:, :]  # B, H*W
+
+        # coverage
+        num_coverage = N
+
+        coverage_list = []
+        for i in range(B):
+            idx_c = torch.randperm(remaining.size()[1])  # shuffles "H*W - beta*N"
+            coverage_list.append(remaining[i, :][idx_c[:num_coverage]].view(1, -1))  # 1, N-beta*N
+        coverage = torch.cat(coverage_list, dim=0)  # B, N-beta*N
+        samples = coverage
+
+    # point coordinates
+    rows_int = samples // W         # 0 for first row, H-1 for last row
+    rows_float = rows_int / float(H-1)         # 0 to 1.0
+    rows_float = (rows_float * 2.0) - 1.0       # -1.0 to 1.0
+
+    cols_int = samples % W          # 0 for first column, W-1 for last column
+    cols_float = cols_int / float(W-1)         # 0 to 1.0
+    cols_float = (cols_float * 2.0) - 1.0       # -1.0 to 1.0
+
+    point_coords = torch.zeros(B, 1, N, 2)
+    point_coords[:, 0, :, 0] = cols_float             # x coord
+    point_coords[:, 0, :, 1] = rows_float             # y coord
+    point_coords = point_coords.to(device)
+    return point_coords, rows_int, cols_int

invokeai/backend/image_util/pbr_maps/architecture/block.py

@@ -0,0 +1,367 @@
+# Original: https://github.com/joeyballentine/Material-Map-Generator
+# Adopted and optimized for Invoke AI
+
+from collections import OrderedDict
+from typing import Any, List, Literal, Optional
+
+import torch
+import torch.nn as nn
+
+ACTIVATION_LAYER_TYPE = Literal["relu", "leakyrelu", "prelu"]
+NORMALIZATION_LAYER_TYPE = Literal["batch", "instance"]
+PADDING_LAYER_TYPE = Literal["zero", "reflect", "replicate"]
+BLOCK_MODE = Literal["CNA", "NAC", "CNAC"]
+UPCONV_BLOCK_MODE = Literal["nearest", "linear", "bilinear", "bicubic", "trilinear"]
+
+
+def act(act_type: ACTIVATION_LAYER_TYPE, inplace: bool = True, neg_slope: float = 0.2, n_prelu: int = 1):
+    """Helper to select Activation Layer"""
+    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)
+    return layer
+
+
+def norm(norm_type: NORMALIZATION_LAYER_TYPE, nc: int):
+    """Helper to select Normalization Layer"""
+    if norm_type == "batch":
+        layer = nn.BatchNorm2d(nc, affine=True)
+    elif norm_type == "instance":
+        layer = nn.InstanceNorm2d(nc, affine=False)
+    return layer
+
+
+def pad(pad_type: PADDING_LAYER_TYPE, padding: int):
+    """Helper to select Padding Layer"""
+    if padding == 0 or pad_type == "zero":
+        return None
+    if pad_type == "reflect":
+        layer = nn.ReflectionPad2d(padding)
+    elif pad_type == "replicate":
+        layer = nn.ReplicationPad2d(padding)
+    return layer
+
+
+def get_valid_padding(kernel_size: int, dilation: int):
+    kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1)
+    padding = (kernel_size - 1) // 2
+    return padding
+
+
+def sequential(*args: Any):
+    # 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: List[nn.Module] = []
+    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)
+
+
+def conv_block(
+    in_nc: int,
+    out_nc: int,
+    kernel_size: int,
+    stride: int = 1,
+    dilation: int = 1,
+    groups: int = 1,
+    bias: bool = True,
+    pad_type: Optional[PADDING_LAYER_TYPE] = "zero",
+    norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+    act_type: Optional[ACTIVATION_LAYER_TYPE] = "relu",
+    mode: BLOCK_MODE = "CNA",
+):
+    """
+    Conv layer with padding, normalization, activation
+    mode: CNA --> Conv -> Norm -> Act
+        NAC --> Norm -> Act --> Conv (Identity Mappings in Deep Residual Networks, ECCV16)
+    """
+    assert mode in ["CNA", "NAC", "CNAC"], f"Wrong conv mode [{mode}]"
+    padding = get_valid_padding(kernel_size, dilation)
+    p = pad(pad_type, padding) if pad_type 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
+    match mode:
+        case "CNA":
+            n = norm(norm_type, out_nc) if norm_type else None
+            return sequential(p, c, n, a)
+        case "NAC":
+            if norm_type is None and act_type is not None:
+                a = act(act_type, inplace=False)
+            n = norm(norm_type, in_nc) if norm_type else None
+            return sequential(n, a, p, c)
+        case "CNAC":
+            n = norm(norm_type, in_nc) if norm_type else None
+            return sequential(n, a, p, c)
+
+
+class ConcatBlock(nn.Module):
+    # Concat the output of a submodule to its input
+    def __init__(self, submodule: nn.Module):
+        super(ConcatBlock, self).__init__()
+        self.sub = submodule
+
+    def forward(self, x: torch.Tensor):
+        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: nn.Module):
+        super(ShortcutBlock, self).__init__()
+        self.sub = submodule
+
+    def forward(self, x: torch.Tensor):
+        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: nn.Module):
+        super(ShortcutBlockSPSR, self).__init__()
+        self.sub = submodule
+
+    def forward(self, x: torch.Tensor):
+        return x, self.sub
+
+    def __repr__(self):
+        tmpstr = "Identity + \n|"
+        modstr = self.sub.__repr__().replace("\n", "\n|")
+        tmpstr = tmpstr + modstr
+        return tmpstr
+
+
+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: int,
+        mid_nc: int,
+        out_nc: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        pad_type: PADDING_LAYER_TYPE = "zero",
+        norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+        act_type: Optional[ACTIVATION_LAYER_TYPE] = "relu",
+        mode: BLOCK_MODE = "CNA",
+        res_scale: int = 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
+        )
+
+        self.res = sequential(conv0, conv1)
+        self.res_scale = res_scale
+
+    def forward(self, x: torch.Tensor):
+        res = self.res(x).mul(self.res_scale)
+        return x + res
+
+
+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)
+    """
+
+    def __init__(
+        self,
+        nc: int,
+        kernel_size: int = 3,
+        gc: int = 32,
+        stride: int = 1,
+        bias: bool = True,
+        pad_type: PADDING_LAYER_TYPE = "zero",
+        norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+        act_type: ACTIVATION_LAYER_TYPE = "leakyrelu",
+        mode: BLOCK_MODE = "CNA",
+    ):
+        super(ResidualDenseBlock_5C, self).__init__()
+        # gc: growth channel, i.e. intermediate channels
+        self.conv1 = conv_block(
+            nc, gc, kernel_size, stride, bias=bias, pad_type=pad_type, norm_type=norm_type, act_type=act_type, mode=mode
+        )
+        self.conv2 = conv_block(
+            nc + gc,
+            gc,
+            kernel_size,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=act_type,
+            mode=mode,
+        )
+        self.conv3 = conv_block(
+            nc + 2 * gc,
+            gc,
+            kernel_size,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=act_type,
+            mode=mode,
+        )
+        self.conv4 = conv_block(
+            nc + 3 * gc,
+            gc,
+            kernel_size,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=act_type,
+            mode=mode,
+        )
+        if mode == "CNA":
+            last_act = None
+        else:
+            last_act = act_type
+        self.conv5 = conv_block(
+            nc + 4 * gc, nc, 3, stride, bias=bias, pad_type=pad_type, norm_type=norm_type, act_type=last_act, mode=mode
+        )
+
+    def forward(self, x: torch.Tensor):
+        x1 = self.conv1(x)
+        x2 = self.conv2(torch.cat((x, x1), 1))
+        x3 = self.conv3(torch.cat((x, x1, x2), 1))
+        x4 = self.conv4(torch.cat((x, x1, x2, x3), 1))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5.mul(0.2) + x
+
+
+class RRDB(nn.Module):
+    """
+    Residual in Residual Dense Block
+    (ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks)
+    """
+
+    def __init__(
+        self,
+        nc: int,
+        kernel_size: int = 3,
+        gc: int = 32,
+        stride: int = 1,
+        bias: bool = True,
+        pad_type: PADDING_LAYER_TYPE = "zero",
+        norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+        act_type: ACTIVATION_LAYER_TYPE = "leakyrelu",
+        mode: BLOCK_MODE = "CNA",
+    ):
+        super(RRDB, self).__init__()
+        self.RDB1 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, norm_type, act_type, mode)
+        self.RDB2 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, norm_type, act_type, mode)
+        self.RDB3 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, norm_type, act_type, mode)
+
+    def forward(self, x: torch.Tensor):
+        out = self.RDB1(x)
+        out = self.RDB2(out)
+        out = self.RDB3(out)
+        return out.mul(0.2) + x
+
+
+# Upsampler
+def pixelshuffle_block(
+    in_nc: int,
+    out_nc: int,
+    upscale_factor: int = 2,
+    kernel_size: int = 3,
+    stride: int = 1,
+    bias: bool = True,
+    pad_type: PADDING_LAYER_TYPE = "zero",
+    norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+    act_type: ACTIVATION_LAYER_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: int = 2,
+    kernel_size: int = 3,
+    stride: int = 1,
+    bias: bool = True,
+    pad_type: PADDING_LAYER_TYPE = "zero",
+    norm_type: Optional[NORMALIZATION_LAYER_TYPE] = None,
+    act_type: ACTIVATION_LAYER_TYPE = "relu",
+    mode: UPCONV_BLOCK_MODE = "nearest",
+):
+    # Adopted from 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
+    )
+    return sequential(upsample, conv)

invokeai/backend/image_util/pbr_maps/architecture/pbr_rrdb_net.py

@@ -0,0 +1,70 @@
+# Original: https://github.com/joeyballentine/Material-Map-Generator
+# Adopted and optimized for Invoke AI
+
+import math
+from typing import Literal, Optional
+
+import torch
+import torch.nn as nn
+
+import invokeai.backend.image_util.pbr_maps.architecture.block as B
+
+UPSCALE_MODE = Literal["upconv", "pixelshuffle"]
+
+
+class PBR_RRDB_Net(nn.Module):
+    def __init__(
+        self,
+        in_nc: int,
+        out_nc: int,
+        nf: int,
+        nb: int,
+        gc: int = 32,
+        upscale: int = 4,
+        norm_type: Optional[B.NORMALIZATION_LAYER_TYPE] = None,
+        act_type: B.ACTIVATION_LAYER_TYPE = "leakyrelu",
+        mode: B.BLOCK_MODE = "CNA",
+        res_scale: int = 1,
+        upsample_mode: UPSCALE_MODE = "upconv",
+    ):
+        super(PBR_RRDB_Net, self).__init__()
+        n_upscale = int(math.log(upscale, 2))
+        if upscale == 3:
+            n_upscale = 1
+
+        fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None)
+        rb_blocks = [
+            B.RRDB(
+                nf,
+                kernel_size=3,
+                gc=32,
+                stride=1,
+                bias=True,
+                pad_type="zero",
+                norm_type=norm_type,
+                act_type=act_type,
+                mode="CNA",
+            )
+            for _ in range(nb)
+        ]
+        LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode)
+
+        if upsample_mode == "upconv":
+            upsample_block = B.upconv_block
+        elif upsample_mode == "pixelshuffle":
+            upsample_block = B.pixelshuffle_block
+
+        if upscale == 3:
+            upsampler = upsample_block(nf, nf, 3, act_type=act_type)
+        else:
+            upsampler = [upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)]
+
+        HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
+        HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)
+
+        self.model = B.sequential(
+            fea_conv, B.ShortcutBlock(B.sequential(*rb_blocks, LR_conv)), *upsampler, HR_conv0, HR_conv1
+        )
+
+    def forward(self, x: torch.Tensor):
+        return self.model(x)

invokeai/backend/image_util/pbr_maps/pbr_maps.py

@@ -0,0 +1,141 @@
+# Original: https://github.com/joeyballentine/Material-Map-Generator
+# Adopted and optimized for Invoke AI
+
+import pathlib
+from typing import Any, Literal
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+import torch
+from PIL import Image
+from safetensors.torch import load_file
+
+from invokeai.backend.image_util.pbr_maps.architecture.pbr_rrdb_net import PBR_RRDB_Net
+from invokeai.backend.image_util.pbr_maps.utils.image_ops import crop_seamless, esrgan_launcher_split_merge
+
+NORMAL_MAP_MODEL = (
+    "https://huggingface.co/InvokeAI/pbr-material-maps/resolve/main/normal_map_generator.safetensors?download=true"
+)
+OTHER_MAP_MODEL = (
+    "https://huggingface.co/InvokeAI/pbr-material-maps/resolve/main/franken_map_generator.safetensors?download=true"
+)
+
+
+class PBRMapsGenerator:
+    def __init__(self, normal_map_model: PBR_RRDB_Net, other_map_model: PBR_RRDB_Net, device: torch.device) -> None:
+        self.normal_map_model = normal_map_model
+        self.other_map_model = other_map_model
+        self.device = device
+
+    @staticmethod
+    def load_model(model_path: pathlib.Path, device: torch.device) -> PBR_RRDB_Net:
+        state_dict = load_file(model_path.as_posix(), device=device.type)
+
+        model = PBR_RRDB_Net(
+            3,
+            3,
+            32,
+            12,
+            gc=32,
+            upscale=1,
+            norm_type=None,
+            act_type="leakyrelu",
+            mode="CNA",
+            res_scale=1,
+            upsample_mode="upconv",
+        )
+
+        model.load_state_dict(state_dict, strict=False)
+
+        del state_dict
+        if torch.cuda.is_available() and device.type == "cuda":
+            torch.cuda.empty_cache()
+
+        model.eval()
+
+        for _, v in model.named_parameters():
+            v.requires_grad = False
+
+        return model.to(device)
+
+    def process(self, img: npt.NDArray[Any], model: PBR_RRDB_Net):
+        img = img.astype(np.float32) / np.iinfo(img.dtype).max
+        img = img[..., ::-1].copy()
+        tensor_img = torch.tensor(img).permute(2, 0, 1).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            output = model(tensor_img).data.squeeze(0).float().cpu().clamp_(0, 1).numpy()
+            output = output[[2, 1, 0], :, :]
+            output = np.transpose(output, (1, 2, 0))
+            output = (output * 255.0).round()
+            return output
+
+    def _cv2_to_pil(self, image: npt.NDArray[Any]):
+        return Image.fromarray(cv2.cvtColor(image.astype(np.uint8), cv2.COLOR_RGB2BGR))
+
+    def generate_maps(
+        self,
+        image: Image.Image,
+        tile_size: int = 512,
+        border_mode: Literal["none", "seamless", "mirror", "replicate"] = "none",
+    ):
+        """
+        Generate PBR texture maps (normal, roughness, and displacement) from an input image.
+        The image can optionally be padded before inference to control how borders are treated,
+        which can help create seamless or edge‑consistent textures.
+
+        Args:
+            image: Source image used to generate the PBR maps.
+            tile_size: Maximum tile size used for tiled inference. If the image is larger than
+                this size in either dimension, it will be split into tiles for processing and
+                then merged.
+
+            border_mode: Strategy for padding the image before inference:
+                - "none": No padding is applied; the image is processed as‑is.
+                - "seamless": Pads the image using wrap‑around tiling
+                  (`cv2.BORDER_WRAP`) to help produce seamless textures.
+                - "mirror": Pads the image by mirroring border pixels
+                  (`cv2.BORDER_REFLECT_101`) to reduce edge artifacts.
+                - "replicate": Pads the image by replicating the edge pixels outward
+                  (`cv2.BORDER_REPLICATE`).
+
+        Returns:
+            A tuple of three PIL Images:
+                - normal_map: RGB normal map generated from the input.
+                - roughness: Single‑channel roughness map extracted from the second model output.
+                - displacement: Single‑channel displacement (height) map extracted from the
+                  second model output.
+        """
+
+        models = [self.normal_map_model, self.other_map_model]
+        np_image = np.array(image).astype(np.uint8)
+
+        match border_mode:
+            case "seamless":
+                np_image = cv2.copyMakeBorder(np_image, 16, 16, 16, 16, cv2.BORDER_WRAP)
+            case "mirror":
+                np_image = cv2.copyMakeBorder(np_image, 16, 16, 16, 16, cv2.BORDER_REFLECT_101)
+            case "replicate":
+                np_image = cv2.copyMakeBorder(np_image, 16, 16, 16, 16, cv2.BORDER_REPLICATE)
+            case "none":
+                pass
+
+        img_height, img_width = np_image.shape[:2]
+
+        # Checking whether to perform tiled inference
+        do_split = img_height > tile_size or img_width > tile_size
+
+        if do_split:
+            rlts = esrgan_launcher_split_merge(np_image, self.process, models, scale_factor=1, tile_size=tile_size)
+        else:
+            rlts = [self.process(np_image, model) for model in models]
+
+        if border_mode != "none":
+            rlts = [crop_seamless(rlt) for rlt in rlts]
+
+        normal_map = self._cv2_to_pil(rlts[0])
+        roughness = self._cv2_to_pil(rlts[1][:, :, 1])
+        displacement = self._cv2_to_pil(rlts[1][:, :, 0])
+
+        return normal_map, roughness, displacement

invokeai/backend/image_util/pbr_maps/utils/image_ops.py

@@ -0,0 +1,93 @@
+# Original: https://github.com/joeyballentine/Material-Map-Generator
+# Adopted and optimized for Invoke AI
+
+import math
+from typing import Any, Callable, List
+
+import numpy as np
+import numpy.typing as npt
+
+from invokeai.backend.image_util.pbr_maps.architecture.pbr_rrdb_net import PBR_RRDB_Net
+
+
+def crop_seamless(img: npt.NDArray[Any]):
+    img_height, img_width = img.shape[:2]
+    y, x = 16, 16
+    h, w = img_height - 32, img_width - 32
+    img = img[y : y + h, x : x + w]
+    return img
+
+
+# from https://github.com/ata4/esrgan-launcher/blob/master/upscale.py
+def esrgan_launcher_split_merge(
+    input_image: npt.NDArray[Any],
+    upscale_function: Callable[[npt.NDArray[Any], PBR_RRDB_Net], npt.NDArray[Any]],
+    models: List[PBR_RRDB_Net],
+    scale_factor: int = 4,
+    tile_size: int = 512,
+    tile_padding: float = 0.125,
+):
+    width, height, depth = input_image.shape
+    output_width = width * scale_factor
+    output_height = height * scale_factor
+    output_shape = (output_width, output_height, depth)
+
+    # start with black image
+    output_images = [np.zeros(output_shape, np.uint8) for _ in range(len(models))]
+
+    tile_padding = math.ceil(tile_size * tile_padding)
+    tile_size = math.ceil(tile_size / scale_factor)
+
+    tiles_x = math.ceil(width / tile_size)
+    tiles_y = math.ceil(height / tile_size)
+
+    for y in range(tiles_y):
+        for x in range(tiles_x):
+            # extract tile from input image
+            ofs_x = x * tile_size
+            ofs_y = y * tile_size
+
+            # input tile area on total image
+            input_start_x = ofs_x
+            input_end_x = min(ofs_x + tile_size, width)
+
+            input_start_y = ofs_y
+            input_end_y = min(ofs_y + tile_size, height)
+
+            # input tile area on total image with padding
+            input_start_x_pad = max(input_start_x - tile_padding, 0)
+            input_end_x_pad = min(input_end_x + tile_padding, width)
+
+            input_start_y_pad = max(input_start_y - tile_padding, 0)
+            input_end_y_pad = min(input_end_y + tile_padding, height)
+
+            # input tile dimensions
+            input_tile_width = input_end_x - input_start_x
+            input_tile_height = input_end_y - input_start_y
+
+            input_tile = input_image[input_start_x_pad:input_end_x_pad, input_start_y_pad:input_end_y_pad]
+
+            for idx, model in enumerate(models):
+                # upscale tile
+                output_tile = upscale_function(input_tile, model)
+
+                # output tile area on total image
+                output_start_x = input_start_x * scale_factor
+                output_end_x = input_end_x * scale_factor
+
+                output_start_y = input_start_y * scale_factor
+                output_end_y = input_end_y * scale_factor
+
+                # output tile area without padding
+                output_start_x_tile = (input_start_x - input_start_x_pad) * scale_factor
+                output_end_x_tile = output_start_x_tile + input_tile_width * scale_factor
+
+                output_start_y_tile = (input_start_y - input_start_y_pad) * scale_factor
+                output_end_y_tile = output_start_y_tile + input_tile_height * scale_factor
+
+                # put tile into output image
+                output_images[idx][output_start_x:output_end_x, output_start_y:output_end_y] = output_tile[
+                    output_start_x_tile:output_end_x_tile, output_start_y_tile:output_end_y_tile
+                ]
+
+    return output_images

invokeai/backend/image_util/pidi/__init__.py

@@ -0,0 +1,80 @@
+# Adapted from https://github.com/huggingface/controlnet_aux
+
+import pathlib
+
+import cv2
+import huggingface_hub
+import numpy as np
+import torch
+from einops import rearrange
+from PIL import Image
+
+from invokeai.backend.image_util.pidi.model import PiDiNet, pidinet
+from invokeai.backend.image_util.util import nms, normalize_image_channel_count, np_to_pil, pil_to_np, safe_step
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+class PIDINetDetector:
+    """Simple wrapper around a PiDiNet model for edge detection."""
+
+    hf_repo_id = "lllyasviel/Annotators"
+    hf_filename = "table5_pidinet.pth"
+
+    @classmethod
+    def get_model_url(cls) -> str:
+        """Get the URL to download the model from the Hugging Face Hub."""
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path) -> PiDiNet:
+        """Load the model from a file."""
+
+        model = pidinet()
+        model.load_state_dict({k.replace("module.", ""): v for k, v in torch.load(model_path)["state_dict"].items()})
+        model.eval()
+        return model
+
+    def __init__(self, model: PiDiNet) -> None:
+        self.model = model
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(
+        self, image: Image.Image, quantize_edges: bool = False, scribble: bool = False, apply_filter: bool = False
+    ) -> Image.Image:
+        """Processes an image and returns the detected edges."""
+
+        device = get_effective_device(self.model)
+
+        np_img = pil_to_np(image)
+        np_img = normalize_image_channel_count(np_img)
+
+        assert np_img.ndim == 3
+
+        bgr_img = np_img[:, :, ::-1].copy()
+
+        with torch.no_grad():
+            image_pidi = torch.from_numpy(bgr_img).float().to(device)
+            image_pidi = image_pidi / 255.0
+            image_pidi = rearrange(image_pidi, "h w c -> 1 c h w")
+            edge = self.model(image_pidi)[-1]
+            edge = edge.cpu().numpy()
+            if apply_filter:
+                edge = edge > 0.5
+            if quantize_edges:
+                edge = safe_step(edge)
+            edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
+
+        detected_map = edge[0, 0]
+
+        if scribble:
+            detected_map = nms(detected_map, 127, 3.0)
+            detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
+            detected_map[detected_map > 4] = 255
+            detected_map[detected_map < 255] = 0
+
+        output_img = np_to_pil(detected_map)
+
+        return output_img

invokeai/backend/image_util/pidi/model.py

@@ -0,0 +1,681 @@
+"""
+Author: Zhuo Su, Wenzhe Liu
+Date: Feb 18, 2021
+"""
+
+import math
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def img2tensor(imgs, bgr2rgb=True, float32=True):
+    """Numpy array to tensor.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Input images.
+        bgr2rgb (bool): Whether to change bgr to rgb.
+        float32 (bool): Whether to change to float32.
+
+    Returns:
+        list[tensor] | tensor: Tensor images. If returned results only have
+            one element, just return tensor.
+    """
+
+    def _totensor(img, bgr2rgb, float32):
+        if img.shape[2] == 3 and bgr2rgb:
+            if img.dtype == 'float64':
+                img = img.astype('float32')
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = torch.from_numpy(img.transpose(2, 0, 1))
+        if float32:
+            img = img.float()
+        return img
+
+    if isinstance(imgs, list):
+        return [_totensor(img, bgr2rgb, float32) for img in imgs]
+    else:
+        return _totensor(imgs, bgr2rgb, float32)
+
+nets = {
+    'baseline': {
+        'layer0':  'cv',
+        'layer1':  'cv',
+        'layer2':  'cv',
+        'layer3':  'cv',
+        'layer4':  'cv',
+        'layer5':  'cv',
+        'layer6':  'cv',
+        'layer7':  'cv',
+        'layer8':  'cv',
+        'layer9':  'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'cv',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+        },
+    'c-v15': {
+        'layer0':  'cd',
+        'layer1':  'cv',
+        'layer2':  'cv',
+        'layer3':  'cv',
+        'layer4':  'cv',
+        'layer5':  'cv',
+        'layer6':  'cv',
+        'layer7':  'cv',
+        'layer8':  'cv',
+        'layer9':  'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'cv',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+        },
+    'a-v15': {
+        'layer0':  'ad',
+        'layer1':  'cv',
+        'layer2':  'cv',
+        'layer3':  'cv',
+        'layer4':  'cv',
+        'layer5':  'cv',
+        'layer6':  'cv',
+        'layer7':  'cv',
+        'layer8':  'cv',
+        'layer9':  'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'cv',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+        },
+    'r-v15': {
+        'layer0':  'rd',
+        'layer1':  'cv',
+        'layer2':  'cv',
+        'layer3':  'cv',
+        'layer4':  'cv',
+        'layer5':  'cv',
+        'layer6':  'cv',
+        'layer7':  'cv',
+        'layer8':  'cv',
+        'layer9':  'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'cv',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+        },
+    'cvvv4': {
+        'layer0':  'cd',
+        'layer1':  'cv',
+        'layer2':  'cv',
+        'layer3':  'cv',
+        'layer4':  'cd',
+        'layer5':  'cv',
+        'layer6':  'cv',
+        'layer7':  'cv',
+        'layer8':  'cd',
+        'layer9':  'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'cd',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+        },
+    'avvv4': {
+        'layer0':  'ad',
+        'layer1':  'cv',
+        'layer2':  'cv',
+        'layer3':  'cv',
+        'layer4':  'ad',
+        'layer5':  'cv',
+        'layer6':  'cv',
+        'layer7':  'cv',
+        'layer8':  'ad',
+        'layer9':  'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'ad',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+        },
+    'rvvv4': {
+        'layer0':  'rd',
+        'layer1':  'cv',
+        'layer2':  'cv',
+        'layer3':  'cv',
+        'layer4':  'rd',
+        'layer5':  'cv',
+        'layer6':  'cv',
+        'layer7':  'cv',
+        'layer8':  'rd',
+        'layer9':  'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'rd',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+        },
+    'cccv4': {
+        'layer0':  'cd',
+        'layer1':  'cd',
+        'layer2':  'cd',
+        'layer3':  'cv',
+        'layer4':  'cd',
+        'layer5':  'cd',
+        'layer6':  'cd',
+        'layer7':  'cv',
+        'layer8':  'cd',
+        'layer9':  'cd',
+        'layer10': 'cd',
+        'layer11': 'cv',
+        'layer12': 'cd',
+        'layer13': 'cd',
+        'layer14': 'cd',
+        'layer15': 'cv',
+        },
+    'aaav4': {
+        'layer0':  'ad',
+        'layer1':  'ad',
+        'layer2':  'ad',
+        'layer3':  'cv',
+        'layer4':  'ad',
+        'layer5':  'ad',
+        'layer6':  'ad',
+        'layer7':  'cv',
+        'layer8':  'ad',
+        'layer9':  'ad',
+        'layer10': 'ad',
+        'layer11': 'cv',
+        'layer12': 'ad',
+        'layer13': 'ad',
+        'layer14': 'ad',
+        'layer15': 'cv',
+        },
+    'rrrv4': {
+        'layer0':  'rd',
+        'layer1':  'rd',
+        'layer2':  'rd',
+        'layer3':  'cv',
+        'layer4':  'rd',
+        'layer5':  'rd',
+        'layer6':  'rd',
+        'layer7':  'cv',
+        'layer8':  'rd',
+        'layer9':  'rd',
+        'layer10': 'rd',
+        'layer11': 'cv',
+        'layer12': 'rd',
+        'layer13': 'rd',
+        'layer14': 'rd',
+        'layer15': 'cv',
+        },
+    'c16': {
+        'layer0':  'cd',
+        'layer1':  'cd',
+        'layer2':  'cd',
+        'layer3':  'cd',
+        'layer4':  'cd',
+        'layer5':  'cd',
+        'layer6':  'cd',
+        'layer7':  'cd',
+        'layer8':  'cd',
+        'layer9':  'cd',
+        'layer10': 'cd',
+        'layer11': 'cd',
+        'layer12': 'cd',
+        'layer13': 'cd',
+        'layer14': 'cd',
+        'layer15': 'cd',
+        },
+    'a16': {
+        'layer0':  'ad',
+        'layer1':  'ad',
+        'layer2':  'ad',
+        'layer3':  'ad',
+        'layer4':  'ad',
+        'layer5':  'ad',
+        'layer6':  'ad',
+        'layer7':  'ad',
+        'layer8':  'ad',
+        'layer9':  'ad',
+        'layer10': 'ad',
+        'layer11': 'ad',
+        'layer12': 'ad',
+        'layer13': 'ad',
+        'layer14': 'ad',
+        'layer15': 'ad',
+        },
+    'r16': {
+        'layer0':  'rd',
+        'layer1':  'rd',
+        'layer2':  'rd',
+        'layer3':  'rd',
+        'layer4':  'rd',
+        'layer5':  'rd',
+        'layer6':  'rd',
+        'layer7':  'rd',
+        'layer8':  'rd',
+        'layer9':  'rd',
+        'layer10': 'rd',
+        'layer11': 'rd',
+        'layer12': 'rd',
+        'layer13': 'rd',
+        'layer14': 'rd',
+        'layer15': 'rd',
+        },
+    'carv4': {
+        'layer0':  'cd',
+        'layer1':  'ad',
+        'layer2':  'rd',
+        'layer3':  'cv',
+        'layer4':  'cd',
+        'layer5':  'ad',
+        'layer6':  'rd',
+        'layer7':  'cv',
+        'layer8':  'cd',
+        'layer9':  'ad',
+        'layer10': 'rd',
+        'layer11': 'cv',
+        'layer12': 'cd',
+        'layer13': 'ad',
+        'layer14': 'rd',
+        'layer15': 'cv',
+        },
+    }
+
+def createConvFunc(op_type):
+    assert op_type in ['cv', 'cd', 'ad', 'rd'], 'unknown op type: %s' % str(op_type)
+    if op_type == 'cv':
+        return F.conv2d
+
+    if op_type == 'cd':
+        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
+            assert dilation in [1, 2], 'dilation for cd_conv should be in 1 or 2'
+            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for cd_conv should be 3x3'
+            assert padding == dilation, 'padding for cd_conv set wrong'
+
+            weights_c = weights.sum(dim=[2, 3], keepdim=True)
+            yc = F.conv2d(x, weights_c, stride=stride, padding=0, groups=groups)
+            y = F.conv2d(x, weights, bias, stride=stride, padding=padding, dilation=dilation, groups=groups)
+            return y - yc
+        return func
+    elif op_type == 'ad':
+        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
+            assert dilation in [1, 2], 'dilation for ad_conv should be in 1 or 2'
+            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for ad_conv should be 3x3'
+            assert padding == dilation, 'padding for ad_conv set wrong'
+
+            shape = weights.shape
+            weights = weights.view(shape[0], shape[1], -1)
+            weights_conv = (weights - weights[:, :, [3, 0, 1, 6, 4, 2, 7, 8, 5]]).view(shape) # clock-wise
+            y = F.conv2d(x, weights_conv, bias, stride=stride, padding=padding, dilation=dilation, groups=groups)
+            return y
+        return func
+    elif op_type == 'rd':
+        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
+            assert dilation in [1, 2], 'dilation for rd_conv should be in 1 or 2'
+            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for rd_conv should be 3x3'
+            padding = 2 * dilation
+
+            shape = weights.shape
+            if weights.is_cuda:
+                buffer = torch.cuda.FloatTensor(shape[0], shape[1], 5 * 5).fill_(0)
+            else:
+                buffer = torch.zeros(shape[0], shape[1], 5 * 5).to(weights.device)
+            weights = weights.view(shape[0], shape[1], -1)
+            buffer[:, :, [0, 2, 4, 10, 14, 20, 22, 24]] = weights[:, :, 1:]
+            buffer[:, :, [6, 7, 8, 11, 13, 16, 17, 18]] = -weights[:, :, 1:]
+            buffer[:, :, 12] = 0
+            buffer = buffer.view(shape[0], shape[1], 5, 5)
+            y = F.conv2d(x, buffer, bias, stride=stride, padding=padding, dilation=dilation, groups=groups)
+            return y
+        return func
+    else:
+        print('impossible to be here unless you force that')
+        return None
+
+class Conv2d(nn.Module):
+    def __init__(self, pdc, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=False):
+        super(Conv2d, self).__init__()
+        if in_channels % groups != 0:
+            raise ValueError('in_channels must be divisible by groups')
+        if out_channels % groups != 0:
+            raise ValueError('out_channels must be divisible by groups')
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.groups = groups
+        self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels // groups, kernel_size, kernel_size))
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+        self.pdc = pdc
+
+    def reset_parameters(self):
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in)
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, input):
+
+        return self.pdc(input, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+class CSAM(nn.Module):
+    """
+    Compact Spatial Attention Module
+    """
+    def __init__(self, channels):
+        super(CSAM, self).__init__()
+
+        mid_channels = 4
+        self.relu1 = nn.ReLU()
+        self.conv1 = nn.Conv2d(channels, mid_channels, kernel_size=1, padding=0)
+        self.conv2 = nn.Conv2d(mid_channels, 1, kernel_size=3, padding=1, bias=False)
+        self.sigmoid = nn.Sigmoid()
+        nn.init.constant_(self.conv1.bias, 0)
+
+    def forward(self, x):
+        y = self.relu1(x)
+        y = self.conv1(y)
+        y = self.conv2(y)
+        y = self.sigmoid(y)
+
+        return x * y
+
+class CDCM(nn.Module):
+    """
+    Compact Dilation Convolution based Module
+    """
+    def __init__(self, in_channels, out_channels):
+        super(CDCM, self).__init__()
+
+        self.relu1 = nn.ReLU()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)
+        self.conv2_1 = nn.Conv2d(out_channels, out_channels, kernel_size=3, dilation=5, padding=5, bias=False)
+        self.conv2_2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, dilation=7, padding=7, bias=False)
+        self.conv2_3 = nn.Conv2d(out_channels, out_channels, kernel_size=3, dilation=9, padding=9, bias=False)
+        self.conv2_4 = nn.Conv2d(out_channels, out_channels, kernel_size=3, dilation=11, padding=11, bias=False)
+        nn.init.constant_(self.conv1.bias, 0)
+
+    def forward(self, x):
+        x = self.relu1(x)
+        x = self.conv1(x)
+        x1 = self.conv2_1(x)
+        x2 = self.conv2_2(x)
+        x3 = self.conv2_3(x)
+        x4 = self.conv2_4(x)
+        return x1 + x2 + x3 + x4
+
+
+class MapReduce(nn.Module):
+    """
+    Reduce feature maps into a single edge map
+    """
+    def __init__(self, channels):
+        super(MapReduce, self).__init__()
+        self.conv = nn.Conv2d(channels, 1, kernel_size=1, padding=0)
+        nn.init.constant_(self.conv.bias, 0)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class PDCBlock(nn.Module):
+    def __init__(self, pdc, inplane, ouplane, stride=1):
+        super(PDCBlock, self).__init__()
+        self.stride=stride
+
+        self.stride=stride
+        if self.stride > 1:
+            self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+            self.shortcut = nn.Conv2d(inplane, ouplane, kernel_size=1, padding=0)
+        self.conv1 = Conv2d(pdc, inplane, inplane, kernel_size=3, padding=1, groups=inplane, bias=False)
+        self.relu2 = nn.ReLU()
+        self.conv2 = nn.Conv2d(inplane, ouplane, kernel_size=1, padding=0, bias=False)
+
+    def forward(self, x):
+        if self.stride > 1:
+            x = self.pool(x)
+        y = self.conv1(x)
+        y = self.relu2(y)
+        y = self.conv2(y)
+        if self.stride > 1:
+            x = self.shortcut(x)
+        y = y + x
+        return y
+
+class PDCBlock_converted(nn.Module):
+    """
+    CPDC, APDC can be converted to vanilla 3x3 convolution
+    RPDC can be converted to vanilla 5x5 convolution
+    """
+    def __init__(self, pdc, inplane, ouplane, stride=1):
+        super(PDCBlock_converted, self).__init__()
+        self.stride=stride
+
+        if self.stride > 1:
+            self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+            self.shortcut = nn.Conv2d(inplane, ouplane, kernel_size=1, padding=0)
+        if pdc == 'rd':
+            self.conv1 = nn.Conv2d(inplane, inplane, kernel_size=5, padding=2, groups=inplane, bias=False)
+        else:
+            self.conv1 = nn.Conv2d(inplane, inplane, kernel_size=3, padding=1, groups=inplane, bias=False)
+        self.relu2 = nn.ReLU()
+        self.conv2 = nn.Conv2d(inplane, ouplane, kernel_size=1, padding=0, bias=False)
+
+    def forward(self, x):
+        if self.stride > 1:
+            x = self.pool(x)
+        y = self.conv1(x)
+        y = self.relu2(y)
+        y = self.conv2(y)
+        if self.stride > 1:
+            x = self.shortcut(x)
+        y = y + x
+        return y
+
+class PiDiNet(nn.Module):
+    def __init__(self, inplane, pdcs, dil=None, sa=False, convert=False):
+        super(PiDiNet, self).__init__()
+        self.sa = sa
+        if dil is not None:
+            assert isinstance(dil, int), 'dil should be an int'
+        self.dil = dil
+
+        self.fuseplanes = []
+
+        self.inplane = inplane
+        if convert:
+            if pdcs[0] == 'rd':
+                init_kernel_size = 5
+                init_padding = 2
+            else:
+                init_kernel_size = 3
+                init_padding = 1
+            self.init_block = nn.Conv2d(3, self.inplane,
+                    kernel_size=init_kernel_size, padding=init_padding, bias=False)
+            block_class = PDCBlock_converted
+        else:
+            self.init_block = Conv2d(pdcs[0], 3, self.inplane, kernel_size=3, padding=1)
+            block_class = PDCBlock
+
+        self.block1_1 = block_class(pdcs[1], self.inplane, self.inplane)
+        self.block1_2 = block_class(pdcs[2], self.inplane, self.inplane)
+        self.block1_3 = block_class(pdcs[3], self.inplane, self.inplane)
+        self.fuseplanes.append(self.inplane) # C
+
+        inplane = self.inplane
+        self.inplane = self.inplane * 2
+        self.block2_1 = block_class(pdcs[4], inplane, self.inplane, stride=2)
+        self.block2_2 = block_class(pdcs[5], self.inplane, self.inplane)
+        self.block2_3 = block_class(pdcs[6], self.inplane, self.inplane)
+        self.block2_4 = block_class(pdcs[7], self.inplane, self.inplane)
+        self.fuseplanes.append(self.inplane) # 2C
+
+        inplane = self.inplane
+        self.inplane = self.inplane * 2
+        self.block3_1 = block_class(pdcs[8], inplane, self.inplane, stride=2)
+        self.block3_2 = block_class(pdcs[9], self.inplane, self.inplane)
+        self.block3_3 = block_class(pdcs[10], self.inplane, self.inplane)
+        self.block3_4 = block_class(pdcs[11], self.inplane, self.inplane)
+        self.fuseplanes.append(self.inplane) # 4C
+
+        self.block4_1 = block_class(pdcs[12], self.inplane, self.inplane, stride=2)
+        self.block4_2 = block_class(pdcs[13], self.inplane, self.inplane)
+        self.block4_3 = block_class(pdcs[14], self.inplane, self.inplane)
+        self.block4_4 = block_class(pdcs[15], self.inplane, self.inplane)
+        self.fuseplanes.append(self.inplane) # 4C
+
+        self.conv_reduces = nn.ModuleList()
+        if self.sa and self.dil is not None:
+            self.attentions = nn.ModuleList()
+            self.dilations = nn.ModuleList()
+            for i in range(4):
+                self.dilations.append(CDCM(self.fuseplanes[i], self.dil))
+                self.attentions.append(CSAM(self.dil))
+                self.conv_reduces.append(MapReduce(self.dil))
+        elif self.sa:
+            self.attentions = nn.ModuleList()
+            for i in range(4):
+                self.attentions.append(CSAM(self.fuseplanes[i]))
+                self.conv_reduces.append(MapReduce(self.fuseplanes[i]))
+        elif self.dil is not None:
+            self.dilations = nn.ModuleList()
+            for i in range(4):
+                self.dilations.append(CDCM(self.fuseplanes[i], self.dil))
+                self.conv_reduces.append(MapReduce(self.dil))
+        else:
+            for i in range(4):
+                self.conv_reduces.append(MapReduce(self.fuseplanes[i]))
+
+        self.classifier = nn.Conv2d(4, 1, kernel_size=1) # has bias
+        nn.init.constant_(self.classifier.weight, 0.25)
+        nn.init.constant_(self.classifier.bias, 0)
+
+        # print('initialization done')
+
+    def get_weights(self):
+        conv_weights = []
+        bn_weights = []
+        relu_weights = []
+        for pname, p in self.named_parameters():
+            if 'bn' in pname:
+                bn_weights.append(p)
+            elif 'relu' in pname:
+                relu_weights.append(p)
+            else:
+                conv_weights.append(p)
+
+        return conv_weights, bn_weights, relu_weights
+
+    def forward(self, x):
+        H, W = x.size()[2:]
+
+        x = self.init_block(x)
+
+        x1 = self.block1_1(x)
+        x1 = self.block1_2(x1)
+        x1 = self.block1_3(x1)
+
+        x2 = self.block2_1(x1)
+        x2 = self.block2_2(x2)
+        x2 = self.block2_3(x2)
+        x2 = self.block2_4(x2)
+
+        x3 = self.block3_1(x2)
+        x3 = self.block3_2(x3)
+        x3 = self.block3_3(x3)
+        x3 = self.block3_4(x3)
+
+        x4 = self.block4_1(x3)
+        x4 = self.block4_2(x4)
+        x4 = self.block4_3(x4)
+        x4 = self.block4_4(x4)
+
+        x_fuses = []
+        if self.sa and self.dil is not None:
+            for i, xi in enumerate([x1, x2, x3, x4]):
+                x_fuses.append(self.attentions[i](self.dilations[i](xi)))
+        elif self.sa:
+            for i, xi in enumerate([x1, x2, x3, x4]):
+                x_fuses.append(self.attentions[i](xi))
+        elif self.dil is not None:
+            for i, xi in enumerate([x1, x2, x3, x4]):
+                x_fuses.append(self.dilations[i](xi))
+        else:
+            x_fuses = [x1, x2, x3, x4]
+
+        e1 = self.conv_reduces[0](x_fuses[0])
+        e1 = F.interpolate(e1, (H, W), mode="bilinear", align_corners=False)
+
+        e2 = self.conv_reduces[1](x_fuses[1])
+        e2 = F.interpolate(e2, (H, W), mode="bilinear", align_corners=False)
+
+        e3 = self.conv_reduces[2](x_fuses[2])
+        e3 = F.interpolate(e3, (H, W), mode="bilinear", align_corners=False)
+
+        e4 = self.conv_reduces[3](x_fuses[3])
+        e4 = F.interpolate(e4, (H, W), mode="bilinear", align_corners=False)
+
+        outputs = [e1, e2, e3, e4]
+
+        output = self.classifier(torch.cat(outputs, dim=1))
+        #if not self.training:
+        #    return torch.sigmoid(output)
+
+        outputs.append(output)
+        outputs = [torch.sigmoid(r) for r in outputs]
+        return outputs
+
+def config_model(model):
+    model_options = list(nets.keys())
+    assert model in model_options, \
+        'unrecognized model, please choose from %s' % str(model_options)
+
+    # print(str(nets[model]))
+
+    pdcs = []
+    for i in range(16):
+        layer_name = 'layer%d' % i
+        op = nets[model][layer_name]
+        pdcs.append(createConvFunc(op))
+
+    return pdcs
+
+def pidinet():
+    pdcs = config_model('carv4')
+    dil = 24 #if args.dil else None
+    return PiDiNet(60, pdcs, dil=dil, sa=True)
+
+
+if __name__ == '__main__':
+    model = pidinet()
+    ckp = torch.load('table5_pidinet.pth')['state_dict']
+    model.load_state_dict({k.replace('module.',''):v for k, v in ckp.items()})
+    im = cv2.imread('examples/test_my/cat_v4.png')
+    im = img2tensor(im).unsqueeze(0)/255.
+    res = model(im)[-1]
+    res = res>0.5
+    res = res.float()
+    res = (res[0,0].cpu().data.numpy()*255.).astype(np.uint8)
+    print(res.shape)
+    cv2.imwrite('edge.png', res)

invokeai/backend/image_util/pngwriter.py

@@ -0,0 +1,118 @@
+"""
+Two helper classes for dealing with PNG images and their path names.
+PngWriter -- Converts Images generated by T2I into PNGs, finds
+             appropriate names for them, and writes prompt metadata
+             into the PNG.
+
+Exports function retrieve_metadata(path)
+"""
+
+import json
+import os
+import re
+
+from PIL import Image, PngImagePlugin
+
+# -------------------image generation utils-----
+
+
+class PngWriter:
+    def __init__(self, outdir):
+        self.outdir = outdir
+        os.makedirs(outdir, exist_ok=True)
+
+    # gives the next unique prefix in outdir
+    def unique_prefix(self):
+        # sort reverse alphabetically until we find max+1
+        dirlist = sorted(os.listdir(self.outdir), reverse=True)
+        # find the first filename that matches our pattern or return 000000.0.png
+        existing_name = next(
+            (f for f in dirlist if re.match(r"^(\d+)\..*\.png", f)),
+            "0000000.0.png",
+        )
+        basecount = int(existing_name.split(".", 1)[0]) + 1
+        return f"{basecount:06}"
+
+    # saves image named _image_ to outdir/name, writing metadata from prompt
+    # returns full path of output
+    def save_image_and_prompt_to_png(self, image, dream_prompt, name, metadata=None, compress_level=6):
+        path = os.path.join(self.outdir, name)
+        info = PngImagePlugin.PngInfo()
+        info.add_text("Dream", dream_prompt)
+        if metadata:
+            info.add_text("sd-metadata", json.dumps(metadata))
+        image.save(path, "PNG", pnginfo=info, compress_level=compress_level)
+        return path
+
+    def retrieve_metadata(self, img_basename):
+        """
+        Given a PNG filename stored in outdir, returns the "sd-metadata"
+        metadata stored there, as a dict
+        """
+        path = os.path.join(self.outdir, img_basename)
+        all_metadata = retrieve_metadata(path)
+        return all_metadata["sd-metadata"]
+
+
+def retrieve_metadata(img_path):
+    """
+    Given a path to a PNG image, returns the "sd-metadata"
+    metadata stored there, as a dict
+    """
+    im = Image.open(img_path)
+    if hasattr(im, "text"):
+        md = im.text.get("sd-metadata", "{}")
+        dream_prompt = im.text.get("Dream", "")
+    else:
+        # When trying to retrieve metadata from images without a 'text' payload, such as JPG images.
+        md = "{}"
+        dream_prompt = ""
+    return {"sd-metadata": json.loads(md), "Dream": dream_prompt}
+
+
+def write_metadata(img_path: str, meta: dict):
+    im = Image.open(img_path)
+    info = PngImagePlugin.PngInfo()
+    info.add_text("sd-metadata", json.dumps(meta))
+    im.save(img_path, "PNG", pnginfo=info)
+
+
+class PromptFormatter:
+    def __init__(self, t2i, opt):
+        self.t2i = t2i
+        self.opt = opt
+
+    # note: the t2i object should provide all these values.
+    # there should be no need to or against opt values
+    def normalize_prompt(self):
+        """Normalize the prompt and switches"""
+        t2i = self.t2i
+        opt = self.opt
+
+        switches = []
+        switches.append(f'"{opt.prompt}"')
+        switches.append(f"-s{opt.steps or t2i.steps}")
+        switches.append(f"-W{opt.width or t2i.width}")
+        switches.append(f"-H{opt.height or t2i.height}")
+        switches.append(f"-C{opt.cfg_scale or t2i.cfg_scale}")
+        switches.append(f"-A{opt.sampler_name or t2i.sampler_name}")
+        # to do: put model name into the t2i object
+        #        switches.append(f'--model{t2i.model_name}')
+        if opt.seamless or t2i.seamless:
+            switches.append("--seamless")
+        if opt.init_img:
+            switches.append(f"-I{opt.init_img}")
+        if opt.fit:
+            switches.append("--fit")
+        if opt.strength and opt.init_img is not None:
+            switches.append(f"-f{opt.strength or t2i.strength}")
+        if opt.gfpgan_strength:
+            switches.append(f"-G{opt.gfpgan_strength}")
+        if opt.upscale:
+            switches.append(f"-U {' '.join([str(u) for u in opt.upscale])}")
+        if opt.variation_amount > 0:
+            switches.append(f"-v{opt.variation_amount}")
+        if opt.with_variations:
+            formatted_variations = ",".join(f"{seed}:{weight}" for seed, weight in opt.with_variations)
+            switches.append(f"-V{formatted_variations}")
+        return " ".join(switches)

invokeai/backend/image_util/realesrgan/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2021, Xintao Wang
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

invokeai/backend/image_util/realesrgan/realesrgan.py

@@ -0,0 +1,272 @@
+import math
+from enum import Enum
+from typing import Any, Optional
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+import torch
+from cv2.typing import MatLike
+from tqdm import tqdm
+
+from invokeai.backend.image_util.basicsr.rrdbnet_arch import RRDBNet
+from invokeai.backend.model_manager.taxonomy import AnyModel
+from invokeai.backend.util.devices import TorchDevice
+
+"""
+Adapted from https://github.com/xinntao/Real-ESRGAN/blob/master/realesrgan/utils.py
+License is BSD3, copied to `LICENSE` in this directory.
+
+The adaptation here has a few changes:
+- Remove print statements, use `tqdm` to show progress
+- Remove unused "outscale" logic, which simply scales the final image to a given factor
+- Remove `dni_weight` logic, which was only used when multiple models were used
+- Remove logic to fetch models from network
+- Add types, rename a few things
+"""
+
+
+class ImageMode(str, Enum):
+    L = "L"
+    RGB = "RGB"
+    RGBA = "RGBA"
+
+
+class RealESRGAN:
+    """A helper class for upsampling images with RealESRGAN.
+
+    Args:
+        scale (int): Upsampling scale factor used in the networks. It is usually 2 or 4.
+        model_path (str): The path to the pretrained model. It can be urls (will first download it automatically).
+        model (nn.Module): The defined network. Default: None.
+        tile (int): As too large images result in the out of GPU memory issue, so this tile option will first crop
+            input images into tiles, and then process each of them. Finally, they will be merged into one image.
+            0 denotes for do not use tile. Default: 0.
+        tile_pad (int): The pad size for each tile, to remove border artifacts. Default: 10.
+        pre_pad (int): Pad the input images to avoid border artifacts. Default: 10.
+        half (float): Whether to use half precision during inference. Default: False.
+    """
+
+    output: torch.Tensor
+
+    def __init__(
+        self,
+        scale: int,
+        loadnet: AnyModel,
+        model: RRDBNet,
+        tile: int = 0,
+        tile_pad: int = 10,
+        pre_pad: int = 10,
+        half: bool = False,
+    ) -> None:
+        self.scale = scale
+        self.tile_size = tile
+        self.tile_pad = tile_pad
+        self.pre_pad = pre_pad
+        self.mod_scale: Optional[int] = None
+        self.half = half
+        self.device = TorchDevice.choose_torch_device()
+
+        # prefer to use params_ema
+        if "params_ema" in loadnet:
+            keyname = "params_ema"
+        else:
+            keyname = "params"
+
+        model.load_state_dict(loadnet[keyname], strict=True)
+        model.eval()
+        self.model = model.to(self.device)
+
+        if self.half:
+            self.model = self.model.half()
+
+    def pre_process(self, img: MatLike) -> None:
+        """Pre-process, such as pre-pad and mod pad, so that the images can be divisible"""
+        img_tensor: torch.Tensor = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
+        self.img = img_tensor.unsqueeze(0).to(self.device)
+        if self.half:
+            self.img = self.img.half()
+
+        # pre_pad
+        if self.pre_pad != 0:
+            self.img = torch.nn.functional.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), "reflect")
+        # mod pad for divisible borders
+        if self.scale == 2:
+            self.mod_scale = 2
+        elif self.scale == 1:
+            self.mod_scale = 4
+        if self.mod_scale is not None:
+            self.mod_pad_h, self.mod_pad_w = 0, 0
+            _, _, h, w = self.img.size()
+            if h % self.mod_scale != 0:
+                self.mod_pad_h = self.mod_scale - h % self.mod_scale
+            if w % self.mod_scale != 0:
+                self.mod_pad_w = self.mod_scale - w % self.mod_scale
+            self.img = torch.nn.functional.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), "reflect")
+
+    def process(self) -> None:
+        # model inference
+        self.output = self.model(self.img)
+
+    def tile_process(self) -> None:
+        """It will first crop input images to tiles, and then process each tile.
+        Finally, all the processed tiles are merged into one images.
+
+        Modified from: https://github.com/ata4/esrgan-launcher
+        """
+        batch, channel, height, width = self.img.shape
+        output_height = height * self.scale
+        output_width = width * self.scale
+        output_shape = (batch, channel, output_height, output_width)
+
+        # start with black image
+        self.output = self.img.new_zeros(output_shape)
+        tiles_x = math.ceil(width / self.tile_size)
+        tiles_y = math.ceil(height / self.tile_size)
+
+        # loop over all tiles
+        total_steps = tiles_y * tiles_x
+        for i in tqdm(range(total_steps), desc="Upscaling"):
+            y = i // tiles_x
+            x = i % tiles_x
+            # extract tile from input image
+            ofs_x = x * self.tile_size
+            ofs_y = y * self.tile_size
+            # input tile area on total image
+            input_start_x = ofs_x
+            input_end_x = min(ofs_x + self.tile_size, width)
+            input_start_y = ofs_y
+            input_end_y = min(ofs_y + self.tile_size, height)
+
+            # input tile area on total image with padding
+            input_start_x_pad = max(input_start_x - self.tile_pad, 0)
+            input_end_x_pad = min(input_end_x + self.tile_pad, width)
+            input_start_y_pad = max(input_start_y - self.tile_pad, 0)
+            input_end_y_pad = min(input_end_y + self.tile_pad, height)
+
+            # input tile dimensions
+            input_tile_width = input_end_x - input_start_x
+            input_tile_height = input_end_y - input_start_y
+            input_tile = self.img[
+                :,
+                :,
+                input_start_y_pad:input_end_y_pad,
+                input_start_x_pad:input_end_x_pad,
+            ]
+
+            # upscale tile
+            with torch.no_grad():
+                output_tile = self.model(input_tile)
+
+            # output tile area on total image
+            output_start_x = input_start_x * self.scale
+            output_end_x = input_end_x * self.scale
+            output_start_y = input_start_y * self.scale
+            output_end_y = input_end_y * self.scale
+
+            # output tile area without padding
+            output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale
+            output_end_x_tile = output_start_x_tile + input_tile_width * self.scale
+            output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale
+            output_end_y_tile = output_start_y_tile + input_tile_height * self.scale
+
+            # put tile into output image
+            self.output[:, :, output_start_y:output_end_y, output_start_x:output_end_x] = output_tile[
+                :,
+                :,
+                output_start_y_tile:output_end_y_tile,
+                output_start_x_tile:output_end_x_tile,
+            ]
+
+    def post_process(self) -> torch.Tensor:
+        # remove extra pad
+        if self.mod_scale is not None:
+            _, _, h, w = self.output.size()
+            self.output = self.output[
+                :,
+                :,
+                0 : h - self.mod_pad_h * self.scale,
+                0 : w - self.mod_pad_w * self.scale,
+            ]
+        # remove prepad
+        if self.pre_pad != 0:
+            _, _, h, w = self.output.size()
+            self.output = self.output[
+                :,
+                :,
+                0 : h - self.pre_pad * self.scale,
+                0 : w - self.pre_pad * self.scale,
+            ]
+        return self.output
+
+    @torch.no_grad()
+    def upscale(self, img: MatLike, esrgan_alpha_upscale: bool = True) -> npt.NDArray[Any]:
+        np_img = img.astype(np.float32)
+        alpha: Optional[np.ndarray] = None
+        if np.max(np_img) > 256:
+            # 16-bit image
+            max_range = 65535
+        else:
+            max_range = 255
+        np_img = np_img / max_range
+        if len(np_img.shape) == 2:
+            # grayscale image
+            img_mode = ImageMode.L
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_GRAY2RGB)
+        elif np_img.shape[2] == 4:
+            # RGBA image with alpha channel
+            img_mode = ImageMode.RGBA
+            alpha = np_img[:, :, 3]
+            np_img = np_img[:, :, 0:3]
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB)
+            if esrgan_alpha_upscale:
+                alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB)
+        else:
+            img_mode = ImageMode.RGB
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB)
+
+        # ------------------- process image (without the alpha channel) ------------------- #
+        self.pre_process(np_img)
+        if self.tile_size > 0:
+            self.tile_process()
+        else:
+            self.process()
+        output_tensor = self.post_process()
+        output_img: npt.NDArray[Any] = output_tensor.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+        output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0))
+        if img_mode is ImageMode.L:
+            output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY)
+
+        # ------------------- process the alpha channel if necessary ------------------- #
+        if img_mode is ImageMode.RGBA:
+            if esrgan_alpha_upscale:
+                assert alpha is not None
+                self.pre_process(alpha)
+                if self.tile_size > 0:
+                    self.tile_process()
+                else:
+                    self.process()
+                output_alpha_tensor = self.post_process()
+                output_alpha: npt.NDArray[Any] = output_alpha_tensor.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+                output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0))
+                output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY)
+            else:  # use the cv2 resize for alpha channel
+                assert alpha is not None
+                h, w = alpha.shape[0:2]
+                output_alpha = cv2.resize(
+                    alpha,
+                    (w * self.scale, h * self.scale),
+                    interpolation=cv2.INTER_LINEAR,
+                )
+
+            # merge the alpha channel
+            output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA)
+            output_img[:, :, 3] = output_alpha
+
+        # ------------------------------ return ------------------------------ #
+        if max_range == 65535:  # 16-bit image
+            output = (output_img * 65535.0).round().astype(np.uint16)
+        else:
+            output = (output_img * 255.0).round().astype(np.uint8)
+
+        return output

invokeai/backend/image_util/safety_checker.py

@@ -0,0 +1,84 @@
+"""
+This module defines a singleton object, "safety_checker" that
+wraps the safety_checker model. It respects the global "nsfw_checker"
+configuration variable, that allows the checker to be supressed.
+"""
+
+from pathlib import Path
+
+import numpy as np
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from PIL import Image, ImageFilter
+from transformers import AutoImageProcessor
+
+import invokeai.backend.util.logging as logger
+from invokeai.app.services.config.config_default import get_config
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.silence_warnings import SilenceWarnings
+
+repo_id = "CompVis/stable-diffusion-safety-checker"
+CHECKER_PATH = "core/convert/stable-diffusion-safety-checker"
+
+
+class SafetyChecker:
+    """
+    Wrapper around SafetyChecker model.
+    """
+
+    feature_extractor = None
+    safety_checker = None
+
+    @classmethod
+    def _load_safety_checker(cls):
+        if cls.safety_checker is not None and cls.feature_extractor is not None:
+            return
+
+        try:
+            model_path = get_config().models_path / CHECKER_PATH
+            if model_path.exists():
+                cls.feature_extractor = AutoImageProcessor.from_pretrained(model_path)
+                cls.safety_checker = StableDiffusionSafetyChecker.from_pretrained(model_path)
+            else:
+                model_path.mkdir(parents=True, exist_ok=True)
+                cls.feature_extractor = AutoImageProcessor.from_pretrained(repo_id)
+                cls.feature_extractor.save_pretrained(model_path)
+                cls.safety_checker = StableDiffusionSafetyChecker.from_pretrained(repo_id)
+                cls.safety_checker.save_pretrained(model_path)
+        except Exception as e:
+            logger.warning(f"Could not load NSFW checker: {str(e)}")
+
+    @classmethod
+    def has_nsfw_concept(cls, image: Image.Image) -> bool:
+        cls._load_safety_checker()
+        if cls.safety_checker is None or cls.feature_extractor is None:
+            return False
+        device = TorchDevice.choose_torch_device()
+        features = cls.feature_extractor([image], return_tensors="pt")
+        features.to(device)
+        cls.safety_checker.to(device)
+        x_image = np.array(image).astype(np.float32) / 255.0
+        x_image = x_image[None].transpose(0, 3, 1, 2)
+        with SilenceWarnings():
+            checked_image, has_nsfw_concept = cls.safety_checker(images=x_image, clip_input=features.pixel_values)
+        return has_nsfw_concept[0]
+
+    @classmethod
+    def blur_if_nsfw(cls, image: Image.Image) -> Image.Image:
+        if cls.has_nsfw_concept(image):
+            logger.warning("A potentially NSFW image has been detected. Image will be blurred.")
+            blurry_image = image.filter(filter=ImageFilter.GaussianBlur(radius=32))
+            caution = cls._get_caution_img()
+            # Center the caution image on the blurred image
+            x = (blurry_image.width - caution.width) // 2
+            y = (blurry_image.height - caution.height) // 2
+            blurry_image.paste(caution, (x, y), caution)
+            image = blurry_image
+
+        return image
+
+    @classmethod
+    def _get_caution_img(cls) -> Image.Image:
+        import invokeai.app.assets.images as image_assets
+
+        caution = Image.open(Path(image_assets.__path__[0]) / "caution.png")
+        return caution.resize((caution.width // 2, caution.height // 2))

invokeai/backend/image_util/segment_anything/mask_refinement.py

@@ -0,0 +1,50 @@
+# This file contains utilities for Grounded-SAM mask refinement based on:
+# https://github.com/NielsRogge/Transformers-Tutorials/blob/a39f33ac1557b02ebfb191ea7753e332b5ca933f/Grounding%20DINO/GroundingDINO_with_Segment_Anything.ipynb
+
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+
+
+def mask_to_polygon(mask: npt.NDArray[np.uint8]) -> list[tuple[int, int]]:
+    """Convert a binary mask to a polygon.
+
+    Returns:
+        list[list[int]]: List of (x, y) coordinates representing the vertices of the polygon.
+    """
+    # Find contours in the binary mask.
+    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    # Find the contour with the largest area.
+    largest_contour = max(contours, key=cv2.contourArea)
+
+    # Extract the vertices of the contour.
+    polygon = largest_contour.reshape(-1, 2).tolist()
+
+    return polygon
+
+
+def polygon_to_mask(
+    polygon: list[tuple[int, int]], image_shape: tuple[int, int], fill_value: int = 1
+) -> npt.NDArray[np.uint8]:
+    """Convert a polygon to a segmentation mask.
+
+    Args:
+        polygon (list): List of (x, y) coordinates representing the vertices of the polygon.
+        image_shape (tuple): Shape of the image (height, width) for the mask.
+        fill_value (int): Value to fill the polygon with.
+
+    Returns:
+        np.ndarray: Segmentation mask with the polygon filled (with value 255).
+    """
+    # Create an empty mask.
+    mask = np.zeros(image_shape, dtype=np.uint8)
+
+    # Convert polygon to an array of points.
+    pts = np.array(polygon, dtype=np.int32)
+
+    # Fill the polygon with white color (255).
+    cv2.fillPoly(mask, [pts], color=(fill_value,))
+
+    return mask

invokeai/backend/image_util/segment_anything/segment_anything_2_pipeline.py

@@ -0,0 +1,109 @@
+from typing import Optional
+
+import torch
+from PIL import Image
+
+# Import SAM2 components - these should be available in transformers 4.56.0+
+from transformers.models.sam2 import Sam2Model
+from transformers.models.sam2.processing_sam2 import Sam2Processor
+
+from invokeai.backend.image_util.segment_anything.shared import SAMInput
+from invokeai.backend.raw_model import RawModel
+
+
+class SegmentAnything2Pipeline(RawModel):
+    """A wrapper class for the transformers SAM2 model and processor that makes it compatible with the model manager."""
+
+    def __init__(self, sam2_model: Sam2Model, sam2_processor: Sam2Processor):
+        """Initialize the SAM2 pipeline.
+
+        Args:
+            sam2_model: The SAM2 model
+            sam2_processor: The SAM2 processor (can be Sam2Processor or Sam2VideoProcessor)
+        """
+        self._sam2_model = sam2_model
+        self._sam2_processor = sam2_processor
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        # HACK: The SAM2 pipeline may not work on MPS devices. We only allow it to be moved to CPU or CUDA.
+        if device is not None and device.type not in {"cpu", "cuda"}:
+            device = None
+        self._sam2_model.to(device=device, dtype=dtype)
+
+    def calc_size(self) -> int:
+        # HACK: Fix the circular import issue.
+        from invokeai.backend.model_manager.load.model_util import calc_module_size
+
+        return calc_module_size(self._sam2_model)
+
+    def segment(
+        self,
+        image: Image.Image,
+        inputs: list[SAMInput],
+    ) -> torch.Tensor:
+        """Segment the image using the provided inputs.
+
+        Args:
+            image: The image to segment.
+            inputs: A list of SAMInput objects containing bounding boxes and/or point lists.
+
+        Returns:
+            torch.Tensor: The segmentation masks. dtype: torch.bool. shape: [num_masks, channels, height, width].
+        """
+
+        input_boxes: list[list[float]] = []
+        input_points: list[list[list[float]]] = []
+        input_labels: list[list[int]] = []
+
+        for i in inputs:
+            box: list[float] | None = None
+            points: list[list[float]] | None = None
+            labels: list[int] | None = None
+
+            if i.bounding_box is not None:
+                box: list[float] | None = [
+                    i.bounding_box.x_min,
+                    i.bounding_box.y_min,
+                    i.bounding_box.x_max,
+                    i.bounding_box.y_max,
+                ]
+
+            if i.points is not None:
+                points = []
+                labels = []
+                for point in i.points:
+                    points.append([point.x, point.y])
+                    labels.append(point.label.value)
+
+            if box is not None:
+                input_boxes.append(box)
+            if points is not None:
+                input_points.append(points)
+            if labels is not None:
+                input_labels.append(labels)
+
+        batched_input_boxes = [input_boxes] if input_boxes else None
+        batched_input_points = [input_points] if input_points else None
+        batched_input_labels = [input_labels] if input_labels else None
+
+        processed_inputs = self._sam2_processor(
+            images=image,
+            input_boxes=batched_input_boxes,
+            input_points=batched_input_points,
+            input_labels=batched_input_labels,
+            return_tensors="pt",
+        ).to(self._sam2_model.device)
+
+        # Generate masks using the SAM2 model
+        outputs = self._sam2_model(**processed_inputs)
+
+        # Post-process the masks to get the final segmentation
+        masks = self._sam2_processor.post_process_masks(
+            masks=outputs.pred_masks,
+            original_sizes=processed_inputs.original_sizes,
+            reshaped_input_sizes=processed_inputs.reshaped_input_sizes,
+        )
+
+        # There should be only one batch.
+        assert len(masks) == 1
+        return masks[0]

invokeai/backend/image_util/segment_anything/segment_anything_pipeline.py

@@ -0,0 +1,97 @@
+from typing import Optional
+
+import torch
+from PIL import Image
+from transformers.models.sam import SamModel
+from transformers.models.sam.processing_sam import SamProcessor
+
+from invokeai.backend.image_util.segment_anything.shared import SAMInput
+from invokeai.backend.raw_model import RawModel
+
+
+class SegmentAnythingPipeline(RawModel):
+    """A wrapper class for the transformers SAM model and processor that makes it compatible with the model manager."""
+
+    def __init__(self, sam_model: SamModel, sam_processor: SamProcessor):
+        self._sam_model = sam_model
+        self._sam_processor = sam_processor
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        # HACK(ryand): The SAM pipeline does not work on MPS devices. We only allow it to be moved to CPU or CUDA.
+        if device is not None and device.type not in {"cpu", "cuda"}:
+            device = None
+        self._sam_model.to(device=device, dtype=dtype)
+
+    def calc_size(self) -> int:
+        # HACK(ryand): Fix the circular import issue.
+        from invokeai.backend.model_manager.load.model_util import calc_module_size
+
+        return calc_module_size(self._sam_model)
+
+    def segment(
+        self,
+        image: Image.Image,
+        inputs: list[SAMInput],
+    ) -> torch.Tensor:
+        """Segment the image using the provided inputs.
+
+        Args:
+            image: The image to segment.
+            inputs: A list of SAMInput objects containing bounding boxes and/or point lists.
+
+        Returns:
+            torch.Tensor: The segmentation masks. dtype: torch.bool. shape: [num_masks, channels, height, width].
+        """
+
+        input_boxes: list[list[float]] = []
+        input_points: list[list[list[float]]] = []
+        input_labels: list[list[int]] = []
+
+        for i in inputs:
+            box: list[float] | None = None
+            points: list[list[float]] | None = None
+            labels: list[int] | None = None
+
+            if i.bounding_box is not None:
+                box: list[float] | None = [
+                    i.bounding_box.x_min,
+                    i.bounding_box.y_min,
+                    i.bounding_box.x_max,
+                    i.bounding_box.y_max,
+                ]
+
+            if i.points is not None:
+                points = []
+                labels = []
+                for point in i.points:
+                    points.append([point.x, point.y])
+                    labels.append(point.label.value)
+
+            if box is not None:
+                input_boxes.append(box)
+            if points is not None:
+                input_points.append(points)
+            if labels is not None:
+                input_labels.append(labels)
+
+        batched_input_boxes = [input_boxes] if input_boxes else None
+        batched_input_points = input_points if input_points else None
+        batched_input_labels = input_labels if input_labels else None
+
+        processed_inputs = self._sam_processor(
+            images=image,
+            input_boxes=batched_input_boxes,
+            input_points=batched_input_points,
+            input_labels=batched_input_labels,
+            return_tensors="pt",
+        ).to(self._sam_model.device)
+        outputs = self._sam_model(**processed_inputs)
+        masks = self._sam_processor.post_process_masks(
+            masks=outputs.pred_masks,
+            original_sizes=processed_inputs.original_sizes,
+            reshaped_input_sizes=processed_inputs.reshaped_input_sizes,
+        )
+
+        # There should be only one batch.
+        assert len(masks) == 1
+        return masks[0]