From 014cc4817da6091a978b5a6ad9fef3d7ef2b0271 Mon Sep 17 00:00:00 2001 From: tazlin Date: Thu, 1 Feb 2024 11:57:08 -0500 Subject: [PATCH 1/6] tests: set `TESTS_ONGOING` by default when tests are running --- tests/conftest.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index deac9a3a..cd3cbfe4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,9 @@ import PIL.Image import pytest +os.environ["TESTS_ONGOING"] = "1" + + from hordelib.comfy_horde import Comfy_Horde from hordelib.horde import HordeLib from hordelib.model_manager.hyper import ALL_MODEL_MANAGER_TYPES From 4cc32a96dc7e8f2a57d848a1e55f2526769a39d5 Mon Sep 17 00:00:00 2001 From: tazlin Date: Thu, 1 Feb 2024 11:57:45 -0500 Subject: [PATCH 2/6] feat: use comfyui `c5a369a33ddb622827552716d9b0119035a2e666` --- hordelib/consts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hordelib/consts.py b/hordelib/consts.py index 7604521c..ca6a9fef 100644 --- a/hordelib/consts.py +++ b/hordelib/consts.py @@ -6,7 +6,7 @@ from hordelib.config_path import get_hordelib_path -COMFYUI_VERSION = "b4e915e74560bd2c090f9b4ed6b73b0781b7050e" +COMFYUI_VERSION = "c5a369a33ddb622827552716d9b0119035a2e666" """The exact version of ComfyUI version to load.""" REMOTE_PROXY = "" From 05c65487897c49f2b86f1b393f215b962944c94e Mon Sep 17 00:00:00 2001 From: tazlin Date: Mon, 5 Feb 2024 22:21:49 -0500 Subject: [PATCH 3/6] feat: initial pyinstaller support --- hordelib/__pyinstaller.py | 10 + .../hed/__init__.py | 2 +- .../leres/__init__.py | 5 +- .../leres/leres/net_tools.py | 18 +- .../midas/api.py | 2 +- .../mlsd/__init__.py | 2 +- .../openpose/__init__.py | 2 +- .../pidinet/__init__.py | 2 +- .../uniformer/__init__.py | 10 +- .../uniformer/mmcv/cnn/bricks/activation.py | 18 +- .../uniformer/mmcv/cnn/bricks/conv_module.py | 76 +-- .../uniformer/mmcv/cnn/bricks/drop.py | 11 +- .../uniformer/mmcv/cnn/bricks/norm.py | 87 ++-- .../uniformer/mmcv/cnn/bricks/registry.py | 2 +- .../uniformer/mmcv/cnn/bricks/transformer.py | 352 +++++++------ .../uniformer/mmcv/cnn/utils/flops_counter.py | 133 ++--- .../uniformer/mmcv/cnn/utils/sync_bn.py | 12 +- .../uniformer/mmcv/cnn/utils/weight_init.py | 9 +- .../uniformer/mmcv/engine/test.py | 37 +- .../uniformer/mmcv/fileio/file_client.py | 319 ++++++------ .../uniformer/mmcv/image/io.py | 112 ++-- .../uniformer/mmcv/image/misc.py | 7 +- .../uniformer/mmcv/ops/cc_attention.py | 20 +- .../uniformer/mmcv/ops/deform_conv.py | 231 ++++----- .../mmcv/ops/modulated_deform_conv.py | 156 +++--- .../mmcv/ops/multi_scale_deform_attn.py | 198 ++++--- .../uniformer/mmcv/ops/nms.py | 158 +++--- .../uniformer/mmcv/ops/point_sample.py | 67 +-- .../uniformer/mmcv/ops/points_sampler.py | 47 +- .../uniformer/mmcv/ops/roiaware_pool3d.py | 40 +- .../uniformer/mmcv/ops/saconv.py | 79 ++- .../uniformer/mmcv/ops/sync_bn.py | 157 +++--- .../uniformer/mmcv/ops/upfirdn2d.py | 58 +-- .../uniformer/mmcv/parallel/distributed.py | 57 +- .../mmcv/parallel/distributed_deprecated.py | 33 +- .../uniformer/mmcv/parallel/registry.py | 2 +- .../uniformer/mmcv/runner/base_module.py | 57 +- .../uniformer/mmcv/runner/base_runner.py | 231 ++++----- .../uniformer/mmcv/runner/checkpoint.py | 213 ++++---- .../mmcv/runner/default_constructor.py | 5 +- .../mmcv/runner/epoch_based_runner.py | 91 ++-- .../uniformer/mmcv/runner/fp16_utils.py | 96 ++-- .../uniformer/mmcv/runner/hooks/checkpoint.py | 89 ++-- .../uniformer/mmcv/runner/hooks/evaluation.py | 213 ++++---- .../uniformer/mmcv/runner/hooks/hook.py | 2 +- .../mmcv/runner/hooks/logger/pavi.py | 66 ++- .../mmcv/runner/hooks/logger/tensorboard.py | 26 +- .../mmcv/runner/hooks/logger/text.py | 147 +++--- .../uniformer/mmcv/runner/hooks/lr_updater.py | 239 ++++----- .../mmcv/runner/hooks/momentum_updater.py | 318 +++++------- .../uniformer/mmcv/runner/hooks/optimizer.py | 179 +++---- .../mmcv/runner/iter_based_runner.py | 143 +++--- .../runner/optimizer/default_constructor.py | 128 +++-- .../uniformer/mmcv/runner/utils.py | 19 +- .../uniformer/mmcv/utils/config.py | 284 +++++----- .../uniformer/mmcv/utils/env.py | 58 ++- .../uniformer/mmcv/utils/ext_loader.py | 50 +- .../uniformer/mmcv/utils/timer.py | 11 +- .../uniformer/mmcv/utils/trace.py | 13 +- .../uniformer/mmcv/video/io.py | 72 ++- .../uniformer/mmcv/video/optflow.py | 100 ++-- .../uniformer/mmcv/video/processing.py | 98 ++-- .../uniformer/mmcv/visualization/color.py | 5 +- .../uniformer/mmcv/visualization/image.py | 58 +-- .../uniformer/mmcv/visualization/optflow.py | 15 +- .../uniformer/mmcv_custom/checkpoint.py | 191 +++---- .../uniformer/mmseg/apis/inference.py | 57 +- .../uniformer/mmseg/apis/test.py | 74 +-- .../uniformer/mmseg/apis/train.py | 69 ++- .../mmseg/core/evaluation/class_names.py | 486 ++++++++++++++---- .../mmseg/core/evaluation/eval_hooks.py | 41 +- .../mmseg/core/evaluation/metrics.py | 172 +++---- .../uniformer/mmseg/core/seg/builder.py | 2 +- .../uniformer/mmseg/datasets/builder.py | 11 +- .../uniformer/mmseg/datasets/cityscapes.py | 109 ++-- .../uniformer/mmseg/datasets/custom.py | 128 +++-- .../mmseg/datasets/pipelines/compose.py | 12 +- .../mmseg/datasets/pipelines/formating.py | 61 ++- .../mmseg/datasets/pipelines/loading.py | 71 ++- .../mmseg/datasets/pipelines/test_time_aug.py | 53 +- .../mmseg/datasets/pipelines/transforms.py | 285 +++++----- .../uniformer/mmseg/models/backbones/cgnet.py | 141 +++-- .../mmseg/models/backbones/fast_scnn.py | 224 ++++---- .../uniformer/mmseg/models/backbones/hrnet.py | 278 +++++----- .../mmseg/models/backbones/mobilenet_v2.py | 58 ++- .../mmseg/models/backbones/mobilenet_v3.py | 137 ++--- .../mmseg/models/backbones/resnest.py | 113 ++-- .../mmseg/models/backbones/resnet.py | 271 +++++----- .../mmseg/models/backbones/resnext.py | 57 +- .../uniformer/mmseg/models/backbones/unet.py | 259 +++++----- .../mmseg/models/backbones/uniformer.py | 294 ++++++++--- .../uniformer/mmseg/models/backbones/vit.py | 251 +++++---- .../uniformer/mmseg/models/builder.py | 4 +- .../mmseg/models/decode_heads/ann_head.py | 86 ++-- .../mmseg/models/decode_heads/apc_head.py | 69 +-- .../mmseg/models/decode_heads/aspp_head.py | 29 +- .../mmseg/models/decode_heads/cc_head.py | 5 +- .../mmseg/models/decode_heads/da_head.py | 41 +- .../mmseg/models/decode_heads/decode_head.py | 82 ++- .../mmseg/models/decode_heads/dm_head.py | 46 +- .../mmseg/models/decode_heads/dnl_head.py | 24 +- .../mmseg/models/decode_heads/ema_head.py | 53 +- .../mmseg/models/decode_heads/enc_head.py | 79 ++- .../mmseg/models/decode_heads/fcn_head.py | 20 +- .../mmseg/models/decode_heads/fpn_head.py | 28 +- .../mmseg/models/decode_heads/gc_head.py | 14 +- .../mmseg/models/decode_heads/lraspp_head.py | 55 +- .../mmseg/models/decode_heads/nl_head.py | 11 +- .../mmseg/models/decode_heads/ocr_head.py | 27 +- .../mmseg/models/decode_heads/point_head.py | 131 ++--- .../mmseg/models/decode_heads/psa_head.py | 114 ++-- .../mmseg/models/decode_heads/psp_head.py | 24 +- .../models/decode_heads/sep_aspp_head.py | 49 +- .../mmseg/models/decode_heads/sep_fcn_head.py | 11 +- .../mmseg/models/decode_heads/uper_head.py | 44 +- .../mmseg/models/losses/lovasz_loss.py | 128 ++--- .../uniformer/mmseg/models/losses/utils.py | 17 +- .../uniformer/mmseg/models/necks/fpn.py | 74 ++- .../mmseg/models/necks/multilevel_neck.py | 35 +- .../uniformer/mmseg/models/segmentors/base.py | 58 +-- .../segmentors/cascade_encoder_decoder.py | 49 +- .../models/segmentors/encoder_decoder.py | 89 ++-- .../mmseg/models/utils/inverted_residual.py | 117 +++-- .../uniformer/mmseg/models/utils/res_layer.py | 63 +-- .../uniformer/mmseg/models/utils/se_layer.py | 23 +- .../models/utils/self_attention_block.py | 62 +-- .../mmseg/models/utils/up_conv_block.py | 47 +- .../uniformer/mmseg/utils/collect_env.py | 12 +- .../uniformer/mmseg/utils/logger.py | 4 +- hordelib/nodes/facerestore/__init__.py | 4 +- .../facerestore/facelib/detection/__init__.py | 6 +- .../detection/retinaface/retinaface.py | 6 +- .../detection/yolov5face/face_detector.py | 8 +- .../detection/yolov5face/models/common.py | 87 +--- .../yolov5face/models/experimental.py | 2 +- .../detection/yolov5face/models/yolo.py | 10 +- .../facerestore/facelib/parsing/__init__.py | 2 +- .../facelib/utils/face_restoration_helper.py | 6 +- .../facerestore/facelib/utils/face_utils.py | 4 +- hordelib/pyinstaller_hooks/__init__.py | 0 .../pyinstaller_hooks/hook-hordelib.horde.py | 53 ++ hordelib/utils/distance.py | 6 +- mypy.ini | 2 +- pyinstaller.py | 21 + pyinstaller_test_entrypoint.py | 36 ++ 145 files changed, 5492 insertions(+), 6047 deletions(-) create mode 100644 hordelib/__pyinstaller.py create mode 100644 hordelib/pyinstaller_hooks/__init__.py create mode 100644 hordelib/pyinstaller_hooks/hook-hordelib.horde.py create mode 100644 pyinstaller.py create mode 100644 pyinstaller_test_entrypoint.py diff --git a/hordelib/__pyinstaller.py b/hordelib/__pyinstaller.py new file mode 100644 index 00000000..4c09d225 --- /dev/null +++ b/hordelib/__pyinstaller.py @@ -0,0 +1,10 @@ +import pathlib + + +def get_hook_dirs() -> list[str]: + return [str(pathlib.Path(__file__).parent / "pyinstaller_hooks")] + + +def get_PyInstaller_tests() -> list[str]: + return [] # FIXME + return [str(pathlib.Path(__file__).parent / "pyinstaller_tests")] diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/hed/__init__.py b/hordelib/nodes/comfy_controlnet_preprocessors/hed/__init__.py index a3249f9c..35ea2746 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/hed/__init__.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/hed/__init__.py @@ -115,7 +115,7 @@ def __init__(self): ) modelpath = os.path.join(builtins.annotator_ckpts_path, "network-bsds500.pth") if not os.path.exists(modelpath): - from comfy_controlnet_preprocessors.util import load_file_from_url + from hordelib.nodes.comfy_controlnet_preprocessors.util import load_file_from_url load_file_from_url(remote_model_path, model_dir=builtins.annotator_ckpts_path) self.netNetwork = Network(modelpath).to(model_management.get_torch_device()).eval() diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/leres/__init__.py b/hordelib/nodes/comfy_controlnet_preprocessors/leres/__init__.py index 8453809b..e242becd 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/leres/__init__.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/leres/__init__.py @@ -45,7 +45,7 @@ def download_model_if_not_existed(): # if os.path.exists(old_model_path): # model_path = old_model_path if not os.path.exists(model_path): - from comfy_controlnet_preprocessors.util import load_file_from_url + from hordelib.nodes.comfy_controlnet_preprocessors.util import load_file_from_url load_file_from_url(remote_model_path_leres, model_dir=builtins.annotator_ckpts_path) os.rename(os.path.join(builtins.annotator_ckpts_path, "res101.pth"), model_path) @@ -66,7 +66,7 @@ def apply_leres(input_image, thr_a, thr_b): """ if boost and pix2pixmodel is None: pix2pixmodel_path = os.path.join(builtins.annotator_ckpts_path, "latest_net_G.pth") if not os.path.exists(pix2pixmodel_path): - from comfy_controlnet_preprocessors.util import load_file_from_url + from hordelib.nodes.comfy_controlnet_preprocessors.util import load_file_from_url load_file_from_url(remote_model_path_pix2pix, model_dir=builtins.annotator_ckpts_path) opt = TestOptions().parse() @@ -84,7 +84,6 @@ def apply_leres(input_image, thr_a, thr_b): height, width, dim = input_image.shape with torch.no_grad(): - if boost: depth = estimateboost(input_image, model, 0, pix2pixmodel, max(width, height)) else: diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/leres/leres/net_tools.py b/hordelib/nodes/comfy_controlnet_preprocessors/leres/leres/net_tools.py index 170ef753..2f3f4b41 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/leres/leres/net_tools.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/leres/leres/net_tools.py @@ -9,21 +9,22 @@ def get_func(func_name): function in this module or the path to a function relative to the base 'modeling' module. """ - if func_name == '': + if func_name == "": return None try: - parts = func_name.split('.') + parts = func_name.split(".") # Refers to a function in this module if len(parts) == 1: return globals()[parts[0]] # Otherwise, assume we're referencing a module under modeling - module_name = 'comfy_controlnet_preprocessors.leres.leres.' + '.'.join(parts[:-1]) + module_name = "hordelib.nodes.comfy_controlnet_preprocessors.leres.leres." + ".".join(parts[:-1]) module = importlib.import_module(module_name) return getattr(module, parts[-1]) except Exception: - print('Failed to f1ind function: %s', func_name) + print("Failed to f1ind function: %s", func_name) raise + def load_ckpt(args, depth_model, shift_model, focal_model): """ Load checkpoint. @@ -32,13 +33,10 @@ def load_ckpt(args, depth_model, shift_model, focal_model): print("loading checkpoint %s" % args.load_ckpt) checkpoint = torch.load(args.load_ckpt) if shift_model is not None: - shift_model.load_state_dict(strip_prefix_if_present(checkpoint['shift_model'], 'module.'), - strict=True) + shift_model.load_state_dict(strip_prefix_if_present(checkpoint["shift_model"], "module."), strict=True) if focal_model is not None: - focal_model.load_state_dict(strip_prefix_if_present(checkpoint['focal_model'], 'module.'), - strict=True) - depth_model.load_state_dict(strip_prefix_if_present(checkpoint['depth_model'], "module."), - strict=True) + focal_model.load_state_dict(strip_prefix_if_present(checkpoint["focal_model"], "module."), strict=True) + depth_model.load_state_dict(strip_prefix_if_present(checkpoint["depth_model"], "module."), strict=True) del checkpoint torch.cuda.empty_cache() diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/midas/api.py b/hordelib/nodes/comfy_controlnet_preprocessors/midas/api.py index 5262edc2..ebc87957 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/midas/api.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/midas/api.py @@ -92,7 +92,7 @@ def load_model(model_type): elif model_type == "dpt_hybrid": # DPT-Hybrid if not os.path.exists(model_path): - from comfy_controlnet_preprocessors.util import load_file_from_url + from hordelib.nodes.comfy_controlnet_preprocessors.util import load_file_from_url load_file_from_url(remote_model_path, model_dir=builtins.annotator_ckpts_path) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/mlsd/__init__.py b/hordelib/nodes/comfy_controlnet_preprocessors/mlsd/__init__.py index c8108949..dcd68919 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/mlsd/__init__.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/mlsd/__init__.py @@ -19,7 +19,7 @@ class MLSDdetector: def __init__(self): model_path = os.path.join(builtins.annotator_ckpts_path, "mlsd_large_512_fp32.pth") if not os.path.exists(model_path): - from comfy_controlnet_preprocessors.util import load_file_from_url + from hordelib.nodes.comfy_controlnet_preprocessors.util import load_file_from_url load_file_from_url(remote_model_path, model_dir=builtins.annotator_ckpts_path) model = MobileV2_MLSD_Large() diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/openpose/__init__.py b/hordelib/nodes/comfy_controlnet_preprocessors/openpose/__init__.py index 8cce1fe4..909713a7 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/openpose/__init__.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/openpose/__init__.py @@ -20,7 +20,7 @@ def __init__(self): hand_modelpath = os.path.join(builtins.annotator_ckpts_path, "hand_pose_model.pth") if not os.path.exists(hand_modelpath): - from comfy_controlnet_preprocessors.util import load_file_from_url + from hordelib.nodes.comfy_controlnet_preprocessors.util import load_file_from_url load_file_from_url(body_model_path, model_dir=builtins.annotator_ckpts_path) load_file_from_url(hand_model_path, model_dir=builtins.annotator_ckpts_path) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/pidinet/__init__.py b/hordelib/nodes/comfy_controlnet_preprocessors/pidinet/__init__.py index d554a022..95a31523 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/pidinet/__init__.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/pidinet/__init__.py @@ -3,7 +3,7 @@ import numpy as np from einops import rearrange from .model import pidinet -from comfy_controlnet_preprocessors.util import load_state_dict, load_file_from_url +from hordelib.nodes.comfy_controlnet_preprocessors.util import load_state_dict, load_file_from_url import builtins import model_management diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/__init__.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/__init__.py index 47a9fe59..9ff36934 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/__init__.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/__init__.py @@ -1,7 +1,11 @@ import os -from comfy_controlnet_preprocessors.uniformer.mmseg.apis import init_segmentor, inference_segmentor, show_result_pyplot -from comfy_controlnet_preprocessors.uniformer.mmseg.core.evaluation import get_palette +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.apis import ( + init_segmentor, + inference_segmentor, + show_result_pyplot, +) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.core.evaluation import get_palette import builtins import model_management @@ -14,7 +18,7 @@ class UniformerDetector: def __init__(self): modelpath = os.path.join(builtins.annotator_ckpts_path, "upernet_global_small.pth") if not os.path.exists(modelpath): - from comfy_controlnet_preprocessors.util import load_file_from_url + from hordelib.nodes.comfy_controlnet_preprocessors.util import load_file_from_url load_file_from_url(checkpoint_file, model_dir=builtins.annotator_ckpts_path) config_file = os.path.join(os.path.dirname(__file__), "exp", "upernet_global_small", "config.py") diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/activation.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/activation.py index 45271c35..bc229576 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/activation.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/activation.py @@ -3,17 +3,18 @@ import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, build_from_cfg, digit_version +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import ( + TORCH_VERSION, + build_from_cfg, + digit_version, +) from .registry import ACTIVATION_LAYERS -for module in [ - nn.ReLU, nn.LeakyReLU, nn.PReLU, nn.RReLU, nn.ReLU6, nn.ELU, - nn.Sigmoid, nn.Tanh -]: +for module in [nn.ReLU, nn.LeakyReLU, nn.PReLU, nn.RReLU, nn.ReLU6, nn.ELU, nn.Sigmoid, nn.Tanh]: ACTIVATION_LAYERS.register_module(module=module) -@ACTIVATION_LAYERS.register_module(name='Clip') +@ACTIVATION_LAYERS.register_module(name="Clip") @ACTIVATION_LAYERS.register_module() class Clamp(nn.Module): """Clamp activation layer. @@ -28,7 +29,7 @@ class Clamp(nn.Module): Default to 1. """ - def __init__(self, min=-1., max=1.): + def __init__(self, min=-1.0, max=1.0): super(Clamp, self).__init__() self.min = min self.max = max @@ -71,8 +72,7 @@ def forward(self, input): return F.gelu(input) -if (TORCH_VERSION == 'parrots' - or digit_version(TORCH_VERSION) < digit_version('1.4')): +if TORCH_VERSION == "parrots" or digit_version(TORCH_VERSION) < digit_version("1.4"): ACTIVATION_LAYERS.register_module(module=GELU) else: ACTIVATION_LAYERS.register_module(module=nn.GELU) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/conv_module.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/conv_module.py index 899f3749..34e6bc65 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/conv_module.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/conv_module.py @@ -3,7 +3,7 @@ import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import _BatchNorm, _InstanceNorm +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import _BatchNorm, _InstanceNorm from ..utils import constant_init, kaiming_init from .activation import build_activation_layer from .conv import build_conv_layer @@ -65,29 +65,31 @@ class ConvModule(nn.Module): Default: ('conv', 'norm', 'act'). """ - _abbr_ = 'conv_block' - - def __init__(self, - in_channels, - out_channels, - kernel_size, - stride=1, - padding=0, - dilation=1, - groups=1, - bias='auto', - conv_cfg=None, - norm_cfg=None, - act_cfg=dict(type='ReLU'), - inplace=True, - with_spectral_norm=False, - padding_mode='zeros', - order=('conv', 'norm', 'act')): + _abbr_ = "conv_block" + + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + groups=1, + bias="auto", + conv_cfg=None, + norm_cfg=None, + act_cfg=dict(type="ReLU"), + inplace=True, + with_spectral_norm=False, + padding_mode="zeros", + order=("conv", "norm", "act"), + ): super(ConvModule, self).__init__() assert conv_cfg is None or isinstance(conv_cfg, dict) assert norm_cfg is None or isinstance(norm_cfg, dict) assert act_cfg is None or isinstance(act_cfg, dict) - official_padding_mode = ['zeros', 'circular'] + official_padding_mode = ["zeros", "circular"] self.conv_cfg = conv_cfg self.norm_cfg = norm_cfg self.act_cfg = act_cfg @@ -96,12 +98,12 @@ def __init__(self, self.with_explicit_padding = padding_mode not in official_padding_mode self.order = order assert isinstance(self.order, tuple) and len(self.order) == 3 - assert set(order) == set(['conv', 'norm', 'act']) + assert set(order) == set(["conv", "norm", "act"]) self.with_norm = norm_cfg is not None self.with_activation = act_cfg is not None # if the conv layer is before a norm layer, bias is unnecessary. - if bias == 'auto': + if bias == "auto": bias = not self.with_norm self.with_bias = bias @@ -121,7 +123,8 @@ def __init__(self, padding=conv_padding, dilation=dilation, groups=groups, - bias=bias) + bias=bias, + ) # export the attributes of self.conv to a higher level for convenience self.in_channels = self.conv.in_channels self.out_channels = self.conv.out_channels @@ -139,7 +142,7 @@ def __init__(self, # build normalization layers if self.with_norm: # norm layer is after conv layer - if order.index('norm') > order.index('conv'): + if order.index("norm") > order.index("conv"): norm_channels = out_channels else: norm_channels = in_channels @@ -147,8 +150,7 @@ def __init__(self, self.add_module(self.norm_name, norm) if self.with_bias: if isinstance(norm, (_BatchNorm, _InstanceNorm)): - warnings.warn( - 'Unnecessary conv bias before batch/instance norm') + warnings.warn("Unnecessary conv bias before batch/instance norm") else: self.norm_name = None @@ -156,10 +158,8 @@ def __init__(self, if self.with_activation: act_cfg_ = act_cfg.copy() # nn.Tanh has no 'inplace' argument - if act_cfg_['type'] not in [ - 'Tanh', 'PReLU', 'Sigmoid', 'HSigmoid', 'Swish' - ]: - act_cfg_.setdefault('inplace', inplace) + if act_cfg_["type"] not in ["Tanh", "PReLU", "Sigmoid", "HSigmoid", "Swish"]: + act_cfg_.setdefault("inplace", inplace) self.activate = build_activation_layer(act_cfg_) # Use msra init by default @@ -182,12 +182,12 @@ def init_weights(self): # this method with default ``kaiming_init``. # Note: For PyTorch's conv layers, they will be overwritten by our # initialization implementation using default ``kaiming_init``. - if not hasattr(self.conv, 'init_weights'): - if self.with_activation and self.act_cfg['type'] == 'LeakyReLU': - nonlinearity = 'leaky_relu' - a = self.act_cfg.get('negative_slope', 0.01) + if not hasattr(self.conv, "init_weights"): + if self.with_activation and self.act_cfg["type"] == "LeakyReLU": + nonlinearity = "leaky_relu" + a = self.act_cfg.get("negative_slope", 0.01) else: - nonlinearity = 'relu' + nonlinearity = "relu" a = 0 kaiming_init(self.conv, a=a, nonlinearity=nonlinearity) if self.with_norm: @@ -195,12 +195,12 @@ def init_weights(self): def forward(self, x, activate=True, norm=True): for layer in self.order: - if layer == 'conv': + if layer == "conv": if self.with_explicit_padding: x = self.padding_layer(x) x = self.conv(x) - elif layer == 'norm' and norm and self.with_norm: + elif layer == "norm" and norm and self.with_norm: x = self.norm(x) - elif layer == 'act' and activate and self.with_activation: + elif layer == "act" and activate and self.with_activation: x = self.activate(x) return x diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/drop.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/drop.py index 7e17848e..9411d32b 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/drop.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/drop.py @@ -2,24 +2,23 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv import build_from_cfg +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv import build_from_cfg from .registry import DROPOUT_LAYERS -def drop_path(x, drop_prob=0., training=False): +def drop_path(x, drop_prob=0.0, training=False): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). We follow the implementation https://github.com/rwightman/pytorch-image-models/blob/a2727c1bf78ba0d7b5727f5f95e37fb7f8866b1f/timm/models/layers/drop.py # noqa: E501 """ - if drop_prob == 0. or not training: + if drop_prob == 0.0 or not training: return x keep_prob = 1 - drop_prob # handle tensors with different dimensions, not just 4D tensors. - shape = (x.shape[0], ) + (1, ) * (x.ndim - 1) - random_tensor = keep_prob + torch.rand( - shape, dtype=x.dtype, device=x.device) + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) output = x.div(keep_prob) * random_tensor.floor() return output diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/norm.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/norm.py index 6d84b2a2..780b370c 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/norm.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/norm.py @@ -3,21 +3,25 @@ import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import is_tuple_of -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import SyncBatchNorm, _BatchNorm, _InstanceNorm +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import is_tuple_of +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import ( + SyncBatchNorm, + _BatchNorm, + _InstanceNorm, +) from .registry import NORM_LAYERS -NORM_LAYERS.register_module('BN', module=nn.BatchNorm2d) -NORM_LAYERS.register_module('BN1d', module=nn.BatchNorm1d) -NORM_LAYERS.register_module('BN2d', module=nn.BatchNorm2d) -NORM_LAYERS.register_module('BN3d', module=nn.BatchNorm3d) -NORM_LAYERS.register_module('SyncBN', module=SyncBatchNorm) -NORM_LAYERS.register_module('GN', module=nn.GroupNorm) -NORM_LAYERS.register_module('LN', module=nn.LayerNorm) -NORM_LAYERS.register_module('IN', module=nn.InstanceNorm2d) -NORM_LAYERS.register_module('IN1d', module=nn.InstanceNorm1d) -NORM_LAYERS.register_module('IN2d', module=nn.InstanceNorm2d) -NORM_LAYERS.register_module('IN3d', module=nn.InstanceNorm3d) +NORM_LAYERS.register_module("BN", module=nn.BatchNorm2d) +NORM_LAYERS.register_module("BN1d", module=nn.BatchNorm1d) +NORM_LAYERS.register_module("BN2d", module=nn.BatchNorm2d) +NORM_LAYERS.register_module("BN3d", module=nn.BatchNorm3d) +NORM_LAYERS.register_module("SyncBN", module=SyncBatchNorm) +NORM_LAYERS.register_module("GN", module=nn.GroupNorm) +NORM_LAYERS.register_module("LN", module=nn.LayerNorm) +NORM_LAYERS.register_module("IN", module=nn.InstanceNorm2d) +NORM_LAYERS.register_module("IN1d", module=nn.InstanceNorm1d) +NORM_LAYERS.register_module("IN2d", module=nn.InstanceNorm2d) +NORM_LAYERS.register_module("IN3d", module=nn.InstanceNorm3d) def infer_abbr(class_type): @@ -43,33 +47,32 @@ def infer_abbr(class_type): str: The inferred abbreviation. """ if not inspect.isclass(class_type): - raise TypeError( - f'class_type must be a type, but got {type(class_type)}') - if hasattr(class_type, '_abbr_'): + raise TypeError(f"class_type must be a type, but got {type(class_type)}") + if hasattr(class_type, "_abbr_"): return class_type._abbr_ if issubclass(class_type, _InstanceNorm): # IN is a subclass of BN - return 'in' + return "in" elif issubclass(class_type, _BatchNorm): - return 'bn' + return "bn" elif issubclass(class_type, nn.GroupNorm): - return 'gn' + return "gn" elif issubclass(class_type, nn.LayerNorm): - return 'ln' + return "ln" else: class_name = class_type.__name__.lower() - if 'batch' in class_name: - return 'bn' - elif 'group' in class_name: - return 'gn' - elif 'layer' in class_name: - return 'ln' - elif 'instance' in class_name: - return 'in' + if "batch" in class_name: + return "bn" + elif "group" in class_name: + return "gn" + elif "layer" in class_name: + return "ln" + elif "instance" in class_name: + return "in" else: - return 'norm_layer' + return "norm_layer" -def build_norm_layer(cfg, num_features, postfix=''): +def build_norm_layer(cfg, num_features, postfix=""): """Build normalization layer. Args: @@ -88,14 +91,14 @@ def build_norm_layer(cfg, num_features, postfix=''): created norm layer. """ if not isinstance(cfg, dict): - raise TypeError('cfg must be a dict') - if 'type' not in cfg: + raise TypeError("cfg must be a dict") + if "type" not in cfg: raise KeyError('the cfg dict must contain the key "type"') cfg_ = cfg.copy() - layer_type = cfg_.pop('type') + layer_type = cfg_.pop("type") if layer_type not in NORM_LAYERS: - raise KeyError(f'Unrecognized norm type {layer_type}') + raise KeyError(f"Unrecognized norm type {layer_type}") norm_layer = NORM_LAYERS.get(layer_type) abbr = infer_abbr(norm_layer) @@ -103,14 +106,14 @@ def build_norm_layer(cfg, num_features, postfix=''): assert isinstance(postfix, (int, str)) name = abbr + str(postfix) - requires_grad = cfg_.pop('requires_grad', True) - cfg_.setdefault('eps', 1e-5) - if layer_type != 'GN': + requires_grad = cfg_.pop("requires_grad", True) + cfg_.setdefault("eps", 1e-5) + if layer_type != "GN": layer = norm_layer(num_features, **cfg_) - if layer_type == 'SyncBN' and hasattr(layer, '_specify_ddp_gpu_num'): + if layer_type == "SyncBN" and hasattr(layer, "_specify_ddp_gpu_num"): layer._specify_ddp_gpu_num(1) else: - assert 'num_groups' in cfg_ + assert "num_groups" in cfg_ layer = norm_layer(num_channels=num_features, **cfg_) for param in layer.parameters(): @@ -131,11 +134,11 @@ def is_norm(layer, exclude=None): """ if exclude is not None: if not isinstance(exclude, tuple): - exclude = (exclude, ) + exclude = (exclude,) if not is_tuple_of(exclude, type): raise TypeError( - f'"exclude" must be either None or type or a tuple of types, ' - f'but got {type(exclude)}: {exclude}') + f'"exclude" must be either None or type or a tuple of types, ' f"but got {type(exclude)}: {exclude}" + ) if exclude and isinstance(layer, exclude): return False diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/registry.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/registry.py index e30efe23..bc63fd98 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/registry.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/registry.py @@ -1,6 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. import sys -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry CONV_LAYERS = Registry("conv layer", sys.modules[__name__]) NORM_LAYERS = Registry("norm layer", sys.modules[__name__]) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/transformer.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/transformer.py index e36d22ed..f652c89c 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/transformer.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/bricks/transformer.py @@ -5,29 +5,48 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv import ConfigDict, deprecated_api_warning -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import Linear, build_activation_layer, build_norm_layer -from comfy_controlnet_preprocessors.uniformer.mmcv.runner.base_module import BaseModule, ModuleList, Sequential -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import build_from_cfg +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv import ConfigDict, deprecated_api_warning +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ( + Linear, + build_activation_layer, + build_norm_layer, +) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner.base_module import ( + BaseModule, + ModuleList, + Sequential, +) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import build_from_cfg from .drop import build_dropout -from .registry import (ATTENTION, FEEDFORWARD_NETWORK, POSITIONAL_ENCODING, - TRANSFORMER_LAYER, TRANSFORMER_LAYER_SEQUENCE) +from .registry import ( + ATTENTION, + FEEDFORWARD_NETWORK, + POSITIONAL_ENCODING, + TRANSFORMER_LAYER, + TRANSFORMER_LAYER_SEQUENCE, +) # Avoid BC-breaking of importing MultiScaleDeformableAttention from this file try: - from comfy_controlnet_preprocessors.uniformer.mmcv.ops.multi_scale_deform_attn import MultiScaleDeformableAttention # noqa F401 + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.ops.multi_scale_deform_attn import ( + MultiScaleDeformableAttention, + ) # noqa F401 + warnings.warn( ImportWarning( - '``MultiScaleDeformableAttention`` has been moved to ' - '``mmcv.ops.multi_scale_deform_attn``, please change original path ' # noqa E501 - '``from comfy_controlnet_preprocessors.uniformer.mmcv.cnn.bricks.transformer import MultiScaleDeformableAttention`` ' # noqa E501 - 'to ``from comfy_controlnet_preprocessors.uniformer.mmcv.ops.multi_scale_deform_attn import MultiScaleDeformableAttention`` ' # noqa E501 - )) + "``MultiScaleDeformableAttention`` has been moved to " + "``mmcv.ops.multi_scale_deform_attn``, please change original path " # noqa E501 + "``from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn.bricks.transformer import MultiScaleDeformableAttention`` " # noqa E501 + "to ``from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.ops.multi_scale_deform_attn import MultiScaleDeformableAttention`` " # noqa E501 + ) + ) except ImportError: - warnings.warn('Fail to import ``MultiScaleDeformableAttention`` from ' - '``mmcv.ops.multi_scale_deform_attn``, ' - 'You should install ``mmcv-full`` if you need this module. ') + warnings.warn( + "Fail to import ``MultiScaleDeformableAttention`` from " + "``mmcv.ops.multi_scale_deform_attn``, " + "You should install ``mmcv-full`` if you need this module. " + ) def build_positional_encoding(cfg, default_args=None): @@ -78,47 +97,50 @@ class MultiheadAttention(BaseModule): Default to False. """ - def __init__(self, - embed_dims, - num_heads, - attn_drop=0., - proj_drop=0., - dropout_layer=dict(type='Dropout', drop_prob=0.), - init_cfg=None, - batch_first=False, - **kwargs): + def __init__( + self, + embed_dims, + num_heads, + attn_drop=0.0, + proj_drop=0.0, + dropout_layer=dict(type="Dropout", drop_prob=0.0), + init_cfg=None, + batch_first=False, + **kwargs, + ): super(MultiheadAttention, self).__init__(init_cfg) - if 'dropout' in kwargs: - warnings.warn('The arguments `dropout` in MultiheadAttention ' - 'has been deprecated, now you can separately ' - 'set `attn_drop`(float), proj_drop(float), ' - 'and `dropout_layer`(dict) ') - attn_drop = kwargs['dropout'] - dropout_layer['drop_prob'] = kwargs.pop('dropout') + if "dropout" in kwargs: + warnings.warn( + "The arguments `dropout` in MultiheadAttention " + "has been deprecated, now you can separately " + "set `attn_drop`(float), proj_drop(float), " + "and `dropout_layer`(dict) " + ) + attn_drop = kwargs["dropout"] + dropout_layer["drop_prob"] = kwargs.pop("dropout") self.embed_dims = embed_dims self.num_heads = num_heads self.batch_first = batch_first - self.attn = nn.MultiheadAttention(embed_dims, num_heads, attn_drop, - **kwargs) + self.attn = nn.MultiheadAttention(embed_dims, num_heads, attn_drop, **kwargs) self.proj_drop = nn.Dropout(proj_drop) - self.dropout_layer = build_dropout( - dropout_layer) if dropout_layer else nn.Identity() - - @deprecated_api_warning({'residual': 'identity'}, - cls_name='MultiheadAttention') - def forward(self, - query, - key=None, - value=None, - identity=None, - query_pos=None, - key_pos=None, - attn_mask=None, - key_padding_mask=None, - **kwargs): + self.dropout_layer = build_dropout(dropout_layer) if dropout_layer else nn.Identity() + + @deprecated_api_warning({"residual": "identity"}, cls_name="MultiheadAttention") + def forward( + self, + query, + key=None, + value=None, + identity=None, + query_pos=None, + key_pos=None, + attn_mask=None, + key_padding_mask=None, + **kwargs, + ): """Forward function for `MultiheadAttention`. **kwargs allow passing a more general data flow when combining @@ -171,8 +193,7 @@ def forward(self, if query_pos.shape == key.shape: key_pos = query_pos else: - warnings.warn(f'position encoding of key is' - f'missing in {self.__class__.__name__}.') + warnings.warn(f"position encoding of key is" f"missing in {self.__class__.__name__}.") if query_pos is not None: query = query + query_pos if key_pos is not None: @@ -189,12 +210,7 @@ def forward(self, key = key.transpose(0, 1) value = value.transpose(0, 1) - out = self.attn( - query=query, - key=key, - value=value, - attn_mask=attn_mask, - key_padding_mask=key_padding_mask)[0] + out = self.attn(query=query, key=key, value=value, attn_mask=attn_mask, key_padding_mask=key_padding_mask)[0] if self.batch_first: out = out.transpose(0, 1) @@ -225,25 +241,21 @@ class FFN(BaseModule): Default: None. """ - @deprecated_api_warning( - { - 'dropout': 'ffn_drop', - 'add_residual': 'add_identity' - }, - cls_name='FFN') - def __init__(self, - embed_dims=256, - feedforward_channels=1024, - num_fcs=2, - act_cfg=dict(type='ReLU', inplace=True), - ffn_drop=0., - dropout_layer=None, - add_identity=True, - init_cfg=None, - **kwargs): + @deprecated_api_warning({"dropout": "ffn_drop", "add_residual": "add_identity"}, cls_name="FFN") + def __init__( + self, + embed_dims=256, + feedforward_channels=1024, + num_fcs=2, + act_cfg=dict(type="ReLU", inplace=True), + ffn_drop=0.0, + dropout_layer=None, + add_identity=True, + init_cfg=None, + **kwargs, + ): super(FFN, self).__init__(init_cfg) - assert num_fcs >= 2, 'num_fcs should be no less ' \ - f'than 2. got {num_fcs}.' + assert num_fcs >= 2, "num_fcs should be no less " f"than 2. got {num_fcs}." self.embed_dims = embed_dims self.feedforward_channels = feedforward_channels self.num_fcs = num_fcs @@ -253,19 +265,15 @@ def __init__(self, layers = [] in_channels = embed_dims for _ in range(num_fcs - 1): - layers.append( - Sequential( - Linear(in_channels, feedforward_channels), self.activate, - nn.Dropout(ffn_drop))) + layers.append(Sequential(Linear(in_channels, feedforward_channels), self.activate, nn.Dropout(ffn_drop))) in_channels = feedforward_channels layers.append(Linear(feedforward_channels, embed_dims)) layers.append(nn.Dropout(ffn_drop)) self.layers = Sequential(*layers) - self.dropout_layer = build_dropout( - dropout_layer) if dropout_layer else torch.nn.Identity() + self.dropout_layer = build_dropout(dropout_layer) if dropout_layer else torch.nn.Identity() self.add_identity = add_identity - @deprecated_api_warning({'residual': 'identity'}, cls_name='FFN') + @deprecated_api_warning({"residual": "identity"}, cls_name="FFN") def forward(self, x, identity=None): """Forward function for `FFN`. @@ -316,69 +324,71 @@ class BaseTransformerLayer(BaseModule): or (n, batch, embed_dim). Default to False. """ - def __init__(self, - attn_cfgs=None, - ffn_cfgs=dict( - type='FFN', - embed_dims=256, - feedforward_channels=1024, - num_fcs=2, - ffn_drop=0., - act_cfg=dict(type='ReLU', inplace=True), - ), - operation_order=None, - norm_cfg=dict(type='LN'), - init_cfg=None, - batch_first=False, - **kwargs): - + def __init__( + self, + attn_cfgs=None, + ffn_cfgs=dict( + type="FFN", + embed_dims=256, + feedforward_channels=1024, + num_fcs=2, + ffn_drop=0.0, + act_cfg=dict(type="ReLU", inplace=True), + ), + operation_order=None, + norm_cfg=dict(type="LN"), + init_cfg=None, + batch_first=False, + **kwargs, + ): deprecated_args = dict( - feedforward_channels='feedforward_channels', - ffn_dropout='ffn_drop', - ffn_num_fcs='num_fcs') + feedforward_channels="feedforward_channels", ffn_dropout="ffn_drop", ffn_num_fcs="num_fcs" + ) for ori_name, new_name in deprecated_args.items(): if ori_name in kwargs: warnings.warn( - f'The arguments `{ori_name}` in BaseTransformerLayer ' - f'has been deprecated, now you should set `{new_name}` ' - f'and other FFN related arguments ' - f'to a dict named `ffn_cfgs`. ') + f"The arguments `{ori_name}` in BaseTransformerLayer " + f"has been deprecated, now you should set `{new_name}` " + f"and other FFN related arguments " + f"to a dict named `ffn_cfgs`. " + ) ffn_cfgs[new_name] = kwargs[ori_name] super(BaseTransformerLayer, self).__init__(init_cfg) self.batch_first = batch_first - assert set(operation_order) & set( - ['self_attn', 'norm', 'ffn', 'cross_attn']) == \ - set(operation_order), f'The operation_order of' \ - f' {self.__class__.__name__} should ' \ - f'contains all four operation type ' \ + assert set(operation_order) & set(["self_attn", "norm", "ffn", "cross_attn"]) == set(operation_order), ( + f"The operation_order of" + f" {self.__class__.__name__} should " + f"contains all four operation type " f"{['self_attn', 'norm', 'ffn', 'cross_attn']}" + ) - num_attn = operation_order.count('self_attn') + operation_order.count( - 'cross_attn') + num_attn = operation_order.count("self_attn") + operation_order.count("cross_attn") if isinstance(attn_cfgs, dict): attn_cfgs = [copy.deepcopy(attn_cfgs) for _ in range(num_attn)] else: - assert num_attn == len(attn_cfgs), f'The length ' \ - f'of attn_cfg {num_attn} is ' \ - f'not consistent with the number of attention' \ - f'in operation_order {operation_order}.' + assert num_attn == len(attn_cfgs), ( + f"The length " + f"of attn_cfg {num_attn} is " + f"not consistent with the number of attention" + f"in operation_order {operation_order}." + ) self.num_attn = num_attn self.operation_order = operation_order self.norm_cfg = norm_cfg - self.pre_norm = operation_order[0] == 'norm' + self.pre_norm = operation_order[0] == "norm" self.attentions = ModuleList() index = 0 for operation_name in operation_order: - if operation_name in ['self_attn', 'cross_attn']: - if 'batch_first' in attn_cfgs[index]: - assert self.batch_first == attn_cfgs[index]['batch_first'] + if operation_name in ["self_attn", "cross_attn"]: + if "batch_first" in attn_cfgs[index]: + assert self.batch_first == attn_cfgs[index]["batch_first"] else: - attn_cfgs[index]['batch_first'] = self.batch_first + attn_cfgs[index]["batch_first"] = self.batch_first attention = build_attention(attn_cfgs[index]) # Some custom attentions used as `self_attn` # or `cross_attn` can have different behavior. @@ -389,36 +399,36 @@ def __init__(self, self.embed_dims = self.attentions[0].embed_dims self.ffns = ModuleList() - num_ffns = operation_order.count('ffn') + num_ffns = operation_order.count("ffn") if isinstance(ffn_cfgs, dict): ffn_cfgs = ConfigDict(ffn_cfgs) if isinstance(ffn_cfgs, dict): ffn_cfgs = [copy.deepcopy(ffn_cfgs) for _ in range(num_ffns)] assert len(ffn_cfgs) == num_ffns for ffn_index in range(num_ffns): - if 'embed_dims' not in ffn_cfgs[ffn_index]: - ffn_cfgs['embed_dims'] = self.embed_dims + if "embed_dims" not in ffn_cfgs[ffn_index]: + ffn_cfgs["embed_dims"] = self.embed_dims else: - assert ffn_cfgs[ffn_index]['embed_dims'] == self.embed_dims - self.ffns.append( - build_feedforward_network(ffn_cfgs[ffn_index], - dict(type='FFN'))) + assert ffn_cfgs[ffn_index]["embed_dims"] == self.embed_dims + self.ffns.append(build_feedforward_network(ffn_cfgs[ffn_index], dict(type="FFN"))) self.norms = ModuleList() - num_norms = operation_order.count('norm') + num_norms = operation_order.count("norm") for _ in range(num_norms): self.norms.append(build_norm_layer(norm_cfg, self.embed_dims)[1]) - def forward(self, - query, - key=None, - value=None, - query_pos=None, - key_pos=None, - attn_masks=None, - query_key_padding_mask=None, - key_padding_mask=None, - **kwargs): + def forward( + self, + query, + key=None, + value=None, + query_pos=None, + key_pos=None, + attn_masks=None, + query_key_padding_mask=None, + key_padding_mask=None, + **kwargs, + ): """Forward function for `TransformerDecoderLayer`. **kwargs contains some specific arguments of attentions. @@ -457,19 +467,18 @@ def forward(self, if attn_masks is None: attn_masks = [None for _ in range(self.num_attn)] elif isinstance(attn_masks, torch.Tensor): - attn_masks = [ - copy.deepcopy(attn_masks) for _ in range(self.num_attn) - ] - warnings.warn(f'Use same attn_mask in all attentions in ' - f'{self.__class__.__name__} ') + attn_masks = [copy.deepcopy(attn_masks) for _ in range(self.num_attn)] + warnings.warn(f"Use same attn_mask in all attentions in " f"{self.__class__.__name__} ") else: - assert len(attn_masks) == self.num_attn, f'The length of ' \ - f'attn_masks {len(attn_masks)} must be equal ' \ - f'to the number of attention in ' \ - f'operation_order {self.num_attn}' + assert len(attn_masks) == self.num_attn, ( + f"The length of " + f"attn_masks {len(attn_masks)} must be equal " + f"to the number of attention in " + f"operation_order {self.num_attn}" + ) for layer in self.operation_order: - if layer == 'self_attn': + if layer == "self_attn": temp_key = temp_value = query query = self.attentions[attn_index]( query, @@ -480,15 +489,16 @@ def forward(self, key_pos=query_pos, attn_mask=attn_masks[attn_index], key_padding_mask=query_key_padding_mask, - **kwargs) + **kwargs, + ) attn_index += 1 identity = query - elif layer == 'norm': + elif layer == "norm": query = self.norms[norm_index](query) norm_index += 1 - elif layer == 'cross_attn': + elif layer == "cross_attn": query = self.attentions[attn_index]( query, key, @@ -498,13 +508,13 @@ def forward(self, key_pos=key_pos, attn_mask=attn_masks[attn_index], key_padding_mask=key_padding_mask, - **kwargs) + **kwargs, + ) attn_index += 1 identity = query - elif layer == 'ffn': - query = self.ffns[ffn_index]( - query, identity if self.pre_norm else None) + elif layer == "ffn": + query = self.ffns[ffn_index](query, identity if self.pre_norm else None) ffn_index += 1 return query @@ -533,12 +543,9 @@ class TransformerLayerSequence(BaseModule): def __init__(self, transformerlayers=None, num_layers=None, init_cfg=None): super(TransformerLayerSequence, self).__init__(init_cfg) if isinstance(transformerlayers, dict): - transformerlayers = [ - copy.deepcopy(transformerlayers) for _ in range(num_layers) - ] + transformerlayers = [copy.deepcopy(transformerlayers) for _ in range(num_layers)] else: - assert isinstance(transformerlayers, list) and \ - len(transformerlayers) == num_layers + assert isinstance(transformerlayers, list) and len(transformerlayers) == num_layers self.num_layers = num_layers self.layers = ModuleList() for i in range(num_layers): @@ -546,16 +553,18 @@ def __init__(self, transformerlayers=None, num_layers=None, init_cfg=None): self.embed_dims = self.layers[0].embed_dims self.pre_norm = self.layers[0].pre_norm - def forward(self, - query, - key, - value, - query_pos=None, - key_pos=None, - attn_masks=None, - query_key_padding_mask=None, - key_padding_mask=None, - **kwargs): + def forward( + self, + query, + key, + value, + query_pos=None, + key_pos=None, + attn_masks=None, + query_key_padding_mask=None, + key_padding_mask=None, + **kwargs, + ): """Forward function for `TransformerCoder`. Args: @@ -591,5 +600,6 @@ def forward(self, attn_masks=attn_masks, query_key_padding_mask=query_key_padding_mask, key_padding_mask=key_padding_mask, - **kwargs) + **kwargs, + ) return query diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/flops_counter.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/flops_counter.py index 2690d052..0340ff2b 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/flops_counter.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/flops_counter.py @@ -30,16 +30,12 @@ import torch import torch.nn as nn -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv -def get_model_complexity_info(model, - input_shape, - print_per_layer_stat=True, - as_strings=True, - input_constructor=None, - flush=False, - ost=sys.stdout): +def get_model_complexity_info( + model, input_shape, print_per_layer_stat=True, as_strings=True, input_constructor=None, flush=False, ost=sys.stdout +): """Get complexity information of a model. This method can calculate FLOPs and parameter counts of a model with @@ -95,7 +91,8 @@ def get_model_complexity_info(model, batch = torch.ones(()).new_empty( (1, *input_shape), dtype=next(flops_model.parameters()).dtype, - device=next(flops_model.parameters()).device) + device=next(flops_model.parameters()).device, + ) except StopIteration: # Avoid StopIteration for models which have no parameters, # like `nn.Relu()`, `nn.AvgPool2d`, etc. @@ -105,8 +102,7 @@ def get_model_complexity_info(model, flops_count, params_count = flops_model.compute_average_flops_cost() if print_per_layer_stat: - print_model_with_flops( - flops_model, flops_count, params_count, ost=ost, flush=flush) + print_model_with_flops(flops_model, flops_count, params_count, ost=ost, flush=flush) flops_model.stop_flops_count() if as_strings: @@ -115,7 +111,7 @@ def get_model_complexity_info(model, return flops_count, params_count -def flops_to_string(flops, units='GFLOPs', precision=2): +def flops_to_string(flops, units="GFLOPs", precision=2): """Convert FLOPs number into a string. Note that Here we take a multiply-add counts as one FLOP. @@ -140,22 +136,22 @@ def flops_to_string(flops, units='GFLOPs', precision=2): """ if units is None: if flops // 10**9 > 0: - return str(round(flops / 10.**9, precision)) + ' GFLOPs' + return str(round(flops / 10.0**9, precision)) + " GFLOPs" elif flops // 10**6 > 0: - return str(round(flops / 10.**6, precision)) + ' MFLOPs' + return str(round(flops / 10.0**6, precision)) + " MFLOPs" elif flops // 10**3 > 0: - return str(round(flops / 10.**3, precision)) + ' KFLOPs' + return str(round(flops / 10.0**3, precision)) + " KFLOPs" else: - return str(flops) + ' FLOPs' + return str(flops) + " FLOPs" else: - if units == 'GFLOPs': - return str(round(flops / 10.**9, precision)) + ' ' + units - elif units == 'MFLOPs': - return str(round(flops / 10.**6, precision)) + ' ' + units - elif units == 'KFLOPs': - return str(round(flops / 10.**3, precision)) + ' ' + units + if units == "GFLOPs": + return str(round(flops / 10.0**9, precision)) + " " + units + elif units == "MFLOPs": + return str(round(flops / 10.0**6, precision)) + " " + units + elif units == "KFLOPs": + return str(round(flops / 10.0**3, precision)) + " " + units else: - return str(flops) + ' FLOPs' + return str(flops) + " FLOPs" def params_to_string(num_params, units=None, precision=2): @@ -181,27 +177,21 @@ def params_to_string(num_params, units=None, precision=2): """ if units is None: if num_params // 10**6 > 0: - return str(round(num_params / 10**6, precision)) + ' M' + return str(round(num_params / 10**6, precision)) + " M" elif num_params // 10**3: - return str(round(num_params / 10**3, precision)) + ' k' + return str(round(num_params / 10**3, precision)) + " k" else: return str(num_params) else: - if units == 'M': - return str(round(num_params / 10.**6, precision)) + ' ' + units - elif units == 'K': - return str(round(num_params / 10.**3, precision)) + ' ' + units + if units == "M": + return str(round(num_params / 10.0**6, precision)) + " " + units + elif units == "K": + return str(round(num_params / 10.0**3, precision)) + " " + units else: return str(num_params) -def print_model_with_flops(model, - total_flops, - total_params, - units='GFLOPs', - precision=3, - ost=sys.stdout, - flush=False): +def print_model_with_flops(model, total_flops, total_params, units="GFLOPs", precision=3, ost=sys.stdout, flush=False): """Print a model with FLOPs for each layer. Args: @@ -273,15 +263,15 @@ def accumulate_flops(self): def flops_repr(self): accumulated_num_params = self.accumulate_params() accumulated_flops_cost = self.accumulate_flops() - return ', '.join([ - params_to_string( - accumulated_num_params, units='M', precision=precision), - '{:.3%} Params'.format(accumulated_num_params / total_params), - flops_to_string( - accumulated_flops_cost, units=units, precision=precision), - '{:.3%} FLOPs'.format(accumulated_flops_cost / total_flops), - self.original_extra_repr() - ]) + return ", ".join( + [ + params_to_string(accumulated_num_params, units="M", precision=precision), + "{:.3%} Params".format(accumulated_num_params / total_params), + flops_to_string(accumulated_flops_cost, units=units, precision=precision), + "{:.3%} FLOPs".format(accumulated_flops_cost / total_flops), + self.original_extra_repr(), + ] + ) def add_extra_repr(m): m.accumulate_flops = accumulate_flops.__get__(m) @@ -293,10 +283,10 @@ def add_extra_repr(m): assert m.extra_repr != m.original_extra_repr def del_extra_repr(m): - if hasattr(m, 'original_extra_repr'): + if hasattr(m, "original_extra_repr"): m.extra_repr = m.original_extra_repr del m.original_extra_repr - if hasattr(m, 'accumulate_flops'): + if hasattr(m, "accumulate_flops"): del m.accumulate_flops model.apply(add_extra_repr) @@ -320,14 +310,10 @@ def get_model_parameters_number(model): def add_flops_counting_methods(net_main_module): # adding additional methods to the existing module object, # this is done this way so that each function has access to self object - net_main_module.start_flops_count = start_flops_count.__get__( - net_main_module) - net_main_module.stop_flops_count = stop_flops_count.__get__( - net_main_module) - net_main_module.reset_flops_count = reset_flops_count.__get__( - net_main_module) - net_main_module.compute_average_flops_cost = compute_average_flops_cost.__get__( # noqa: E501 - net_main_module) + net_main_module.start_flops_count = start_flops_count.__get__(net_main_module) + net_main_module.stop_flops_count = stop_flops_count.__get__(net_main_module) + net_main_module.reset_flops_count = reset_flops_count.__get__(net_main_module) + net_main_module.compute_average_flops_cost = compute_average_flops_cost.__get__(net_main_module) # noqa: E501 net_main_module.reset_flops_count() @@ -363,12 +349,11 @@ def start_flops_count(self): def add_flops_counter_hook_function(module): if is_supported_instance(module): - if hasattr(module, '__flops_handle__'): + if hasattr(module, "__flops_handle__"): return else: - handle = module.register_forward_hook( - get_modules_mapping()[type(module)]) + handle = module.register_forward_hook(get_modules_mapping()[type(module)]) module.__flops_handle__ = handle @@ -417,8 +402,7 @@ def relu_flops_counter_hook(module, input, output): def linear_flops_counter_hook(module, input, output): input = input[0] - output_last_dim = output.shape[ - -1] # pytorch checks dimensions, so here we don't care much + output_last_dim = output.shape[-1] # pytorch checks dimensions, so here we don't care much module.__flops__ += int(np.prod(input.shape) * output_last_dim) @@ -431,8 +415,7 @@ def norm_flops_counter_hook(module, input, output): input = input[0] batch_flops = np.prod(input.shape) - if (getattr(module, 'affine', False) - or getattr(module, 'elementwise_affine', False)): + if getattr(module, "affine", False) or getattr(module, "elementwise_affine", False): batch_flops *= 2 module.__flops__ += int(batch_flops) @@ -450,8 +433,7 @@ def deconv_flops_counter_hook(conv_module, input, output): groups = conv_module.groups filters_per_channel = out_channels // groups - conv_per_position_flops = ( - kernel_height * kernel_width * in_channels * filters_per_channel) + conv_per_position_flops = kernel_height * kernel_width * in_channels * filters_per_channel active_elements_count = batch_size * input_height * input_width overall_conv_flops = conv_per_position_flops * active_elements_count @@ -477,8 +459,7 @@ def conv_flops_counter_hook(conv_module, input, output): groups = conv_module.groups filters_per_channel = out_channels // groups - conv_per_position_flops = int( - np.prod(kernel_dims)) * in_channels * filters_per_channel + conv_per_position_flops = int(np.prod(kernel_dims)) * in_channels * filters_per_channel active_elements_count = batch_size * int(np.prod(output_dims)) @@ -487,7 +468,6 @@ def conv_flops_counter_hook(conv_module, input, output): bias_flops = 0 if conv_module.bias is not None: - bias_flops = out_channels * active_elements_count overall_flops = overall_conv_flops + bias_flops @@ -503,18 +483,16 @@ def batch_counter_hook(module, input, output): batch_size = len(input) else: pass - print('Warning! No positional inputs found for a module, ' - 'assuming batch size is 1.') + print("Warning! No positional inputs found for a module, " "assuming batch size is 1.") module.__batch_counter__ += batch_size def add_batch_counter_variables_or_reset(module): - module.__batch_counter__ = 0 def add_batch_counter_hook_function(module): - if hasattr(module, '__batch_counter_handle__'): + if hasattr(module, "__batch_counter_handle__"): return handle = module.register_forward_hook(batch_counter_hook) @@ -522,17 +500,18 @@ def add_batch_counter_hook_function(module): def remove_batch_counter_hook_function(module): - if hasattr(module, '__batch_counter_handle__'): + if hasattr(module, "__batch_counter_handle__"): module.__batch_counter_handle__.remove() del module.__batch_counter_handle__ def add_flops_counter_variable_or_reset(module): if is_supported_instance(module): - if hasattr(module, '__flops__') or hasattr(module, '__params__'): - print('Warning: variables __flops__ or __params__ are already ' - 'defined for the module' + type(module).__name__ + - ' ptflops can affect your code!') + if hasattr(module, "__flops__") or hasattr(module, "__params__"): + print( + "Warning: variables __flops__ or __params__ are already " + "defined for the module" + type(module).__name__ + " ptflops can affect your code!" + ) module.__flops__ = 0 module.__params__ = get_model_parameters_number(module) @@ -545,7 +524,7 @@ def is_supported_instance(module): def remove_flops_counter_hook_function(module): if is_supported_instance(module): - if hasattr(module, '__flops_handle__'): + if hasattr(module, "__flops_handle__"): module.__flops_handle__.remove() del module.__flops_handle__ diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/sync_bn.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/sync_bn.py index 1ea43ac3..051319ce 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/sync_bn.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/sync_bn.py @@ -1,6 +1,6 @@ import torch -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv class _BatchNormXd(torch.nn.modules.batchnorm._BatchNorm): @@ -34,12 +34,12 @@ def revert_sync_batchnorm(module): """ module_output = module module_checklist = [torch.nn.modules.batchnorm.SyncBatchNorm] - if hasattr(mmcv, 'ops'): + if hasattr(mmcv, "ops"): module_checklist.append(mmcv.ops.SyncBatchNorm) if isinstance(module, tuple(module_checklist)): - module_output = _BatchNormXd(module.num_features, module.eps, - module.momentum, module.affine, - module.track_running_stats) + module_output = _BatchNormXd( + module.num_features, module.eps, module.momentum, module.affine, module.track_running_stats + ) if module.affine: # no_grad() may not be needed here but # just to be consistent with `convert_sync_batchnorm()` @@ -51,7 +51,7 @@ def revert_sync_batchnorm(module): module_output.num_batches_tracked = module.num_batches_tracked module_output.training = module.training # qconfig exists in quantized models - if hasattr(module, 'qconfig'): + if hasattr(module, "qconfig"): module_output.qconfig = module.qconfig for name, child in module.named_children(): module_output.add_module(name, revert_sync_batchnorm(child)) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/weight_init.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/weight_init.py index 6ff72d45..fc04fde9 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/weight_init.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/cnn/utils/weight_init.py @@ -9,7 +9,12 @@ import torch.nn as nn from torch import Tensor -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry, build_from_cfg, get_logger, print_log +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import ( + Registry, + build_from_cfg, + get_logger, + print_log, +) INITIALIZERS = Registry("initializer", sys.modules[__name__]) @@ -443,7 +448,7 @@ def __init__(self, checkpoint, prefix=None, map_location=None): self.map_location = map_location def __call__(self, module): - from comfy_controlnet_preprocessors.uniformer.mmcv.runner import ( + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import ( _load_checkpoint_with_prefix, load_checkpoint, load_state_dict, diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/engine/test.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/engine/test.py index 9e6fd610..84613045 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/engine/test.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/engine/test.py @@ -8,8 +8,8 @@ import torch import torch.distributed as dist -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import get_dist_info +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import get_dist_info def single_gpu_test(model, data_loader): @@ -111,22 +111,18 @@ def collect_results_cpu(result_part, size, tmpdir=None): if tmpdir is None: MAX_LEN = 512 # 32 is whitespace - dir_tensor = torch.full((MAX_LEN, ), - 32, - dtype=torch.uint8, - device='cuda') + dir_tensor = torch.full((MAX_LEN,), 32, dtype=torch.uint8, device="cuda") if rank == 0: - mmcv.mkdir_or_exist('.dist_test') - tmpdir = tempfile.mkdtemp(dir='.dist_test') - tmpdir = torch.tensor( - bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda') - dir_tensor[:len(tmpdir)] = tmpdir + mmcv.mkdir_or_exist(".dist_test") + tmpdir = tempfile.mkdtemp(dir=".dist_test") + tmpdir = torch.tensor(bytearray(tmpdir.encode()), dtype=torch.uint8, device="cuda") + dir_tensor[: len(tmpdir)] = tmpdir dist.broadcast(dir_tensor, 0) tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip() else: mmcv.mkdir_or_exist(tmpdir) # dump the part result to the dir - mmcv.dump(result_part, osp.join(tmpdir, f'part_{rank}.pkl')) + mmcv.dump(result_part, osp.join(tmpdir, f"part_{rank}.pkl")) dist.barrier() # collect all parts if rank != 0: @@ -135,7 +131,7 @@ def collect_results_cpu(result_part, size, tmpdir=None): # load results of all parts from tmp dir part_list = [] for i in range(world_size): - part_file = osp.join(tmpdir, f'part_{i}.pkl') + part_file = osp.join(tmpdir, f"part_{i}.pkl") part_result = mmcv.load(part_file) # When data is severely insufficient, an empty part_result # on a certain gpu could makes the overall outputs empty. @@ -169,26 +165,23 @@ def collect_results_gpu(result_part, size): """ rank, world_size = get_dist_info() # dump result part to tensor with pickle - part_tensor = torch.tensor( - bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda') + part_tensor = torch.tensor(bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device="cuda") # gather all result part tensor shape - shape_tensor = torch.tensor(part_tensor.shape, device='cuda') + shape_tensor = torch.tensor(part_tensor.shape, device="cuda") shape_list = [shape_tensor.clone() for _ in range(world_size)] dist.all_gather(shape_list, shape_tensor) # padding result part tensor to max length shape_max = torch.tensor(shape_list).max() - part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda') - part_send[:shape_tensor[0]] = part_tensor - part_recv_list = [ - part_tensor.new_zeros(shape_max) for _ in range(world_size) - ] + part_send = torch.zeros(shape_max, dtype=torch.uint8, device="cuda") + part_send[: shape_tensor[0]] = part_tensor + part_recv_list = [part_tensor.new_zeros(shape_max) for _ in range(world_size)] # gather all result part dist.all_gather(part_recv_list, part_send) if rank == 0: part_list = [] for recv, shape in zip(part_recv_list, shape_list): - part_result = pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()) + part_result = pickle.loads(recv[: shape[0]].cpu().numpy().tobytes()) # When data is severely insufficient, an empty part_result # on a certain gpu could makes the overall outputs empty. if part_result: diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/fileio/file_client.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/fileio/file_client.py index ef4a688e..1fe87036 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/fileio/file_client.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/fileio/file_client.py @@ -11,9 +11,9 @@ from typing import Iterable, Iterator, Optional, Tuple, Union from urllib.request import urlopen -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.misc import has_method -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.path import is_filepath +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.misc import has_method +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.path import is_filepath class BaseStorageBackend(metaclass=ABCMeta): @@ -61,10 +61,9 @@ def __init__(self, path_mapping=None): try: import ceph except ImportError: - raise ImportError('Please install ceph to enable CephBackend.') + raise ImportError("Please install ceph to enable CephBackend.") - warnings.warn( - 'CephBackend will be deprecated, please use PetrelBackend instead') + warnings.warn("CephBackend will be deprecated, please use PetrelBackend instead") self._client = ceph.S3Client() assert isinstance(path_mapping, dict) or path_mapping is None self.path_mapping = path_mapping @@ -105,14 +104,11 @@ class PetrelBackend(BaseStorageBackend): >>> client.get(filepath2) # get data from 'cluster-name' cluster """ - def __init__(self, - path_mapping: Optional[dict] = None, - enable_mc: bool = True): + def __init__(self, path_mapping: Optional[dict] = None, enable_mc: bool = True): try: from petrel_client import client except ImportError: - raise ImportError('Please install petrel_client to enable ' - 'PetrelBackend.') + raise ImportError("Please install petrel_client to enable " "PetrelBackend.") self._client = client.Client(enable_mc=enable_mc) assert isinstance(path_mapping, dict) or path_mapping is None @@ -142,7 +138,7 @@ def _format_path(self, filepath: str) -> str: Args: filepath (str): Path to be formatted. """ - return re.sub(r'\\+', '/', filepath) + return re.sub(r"\\+", "/", filepath) def get(self, filepath: Union[str, Path]) -> memoryview: """Read data from a given ``filepath`` with 'rb' mode. @@ -161,9 +157,7 @@ def get(self, filepath: Union[str, Path]) -> memoryview: value_buf = memoryview(value) return value_buf - def get_text(self, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> str: + def get_text(self, filepath: Union[str, Path], encoding: str = "utf-8") -> str: """Read data from a given ``filepath`` with 'r' mode. Args: @@ -187,10 +181,7 @@ def put(self, obj: bytes, filepath: Union[str, Path]) -> None: filepath = self._format_path(filepath) self._client.put(filepath, obj) - def put_text(self, - obj: str, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> None: + def put_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None: """Save data to a given ``filepath``. Args: @@ -207,11 +198,14 @@ def remove(self, filepath: Union[str, Path]) -> None: Args: filepath (str or Path): Path to be removed. """ - if not has_method(self._client, 'delete'): + if not has_method(self._client, "delete"): raise NotImplementedError( - ('Current version of Petrel Python SDK has not supported ' - 'the `delete` method, please use a higher version or dev' - ' branch instead.')) + ( + "Current version of Petrel Python SDK has not supported " + "the `delete` method, please use a higher version or dev" + " branch instead." + ) + ) filepath = self._map_path(filepath) filepath = self._format_path(filepath) @@ -226,12 +220,14 @@ def exists(self, filepath: Union[str, Path]) -> bool: Returns: bool: Return ``True`` if ``filepath`` exists, ``False`` otherwise. """ - if not (has_method(self._client, 'contains') - and has_method(self._client, 'isdir')): + if not (has_method(self._client, "contains") and has_method(self._client, "isdir")): raise NotImplementedError( - ('Current version of Petrel Python SDK has not supported ' - 'the `contains` and `isdir` methods, please use a higher' - 'version or dev branch instead.')) + ( + "Current version of Petrel Python SDK has not supported " + "the `contains` and `isdir` methods, please use a higher" + "version or dev branch instead." + ) + ) filepath = self._map_path(filepath) filepath = self._format_path(filepath) @@ -248,11 +244,14 @@ def isdir(self, filepath: Union[str, Path]) -> bool: bool: Return ``True`` if ``filepath`` points to a directory, ``False`` otherwise. """ - if not has_method(self._client, 'isdir'): + if not has_method(self._client, "isdir"): raise NotImplementedError( - ('Current version of Petrel Python SDK has not supported ' - 'the `isdir` method, please use a higher version or dev' - ' branch instead.')) + ( + "Current version of Petrel Python SDK has not supported " + "the `isdir` method, please use a higher version or dev" + " branch instead." + ) + ) filepath = self._map_path(filepath) filepath = self._format_path(filepath) @@ -268,18 +267,20 @@ def isfile(self, filepath: Union[str, Path]) -> bool: bool: Return ``True`` if ``filepath`` points to a file, ``False`` otherwise. """ - if not has_method(self._client, 'contains'): + if not has_method(self._client, "contains"): raise NotImplementedError( - ('Current version of Petrel Python SDK has not supported ' - 'the `contains` method, please use a higher version or ' - 'dev branch instead.')) + ( + "Current version of Petrel Python SDK has not supported " + "the `contains` method, please use a higher version or " + "dev branch instead." + ) + ) filepath = self._map_path(filepath) filepath = self._format_path(filepath) return self._client.contains(filepath) - def join_path(self, filepath: Union[str, Path], - *filepaths: Union[str, Path]) -> str: + def join_path(self, filepath: Union[str, Path], *filepaths: Union[str, Path]) -> str: """Concatenate all file paths. Args: @@ -289,12 +290,12 @@ def join_path(self, filepath: Union[str, Path], str: The result after concatenation. """ filepath = self._format_path(self._map_path(filepath)) - if filepath.endswith('/'): + if filepath.endswith("/"): filepath = filepath[:-1] formatted_paths = [filepath] for path in filepaths: formatted_paths.append(self._format_path(self._map_path(path))) - return '/'.join(formatted_paths) + return "/".join(formatted_paths) @contextmanager def get_local_path(self, filepath: Union[str, Path]) -> Iterable[str]: @@ -328,12 +329,14 @@ def get_local_path(self, filepath: Union[str, Path]) -> Iterable[str]: finally: os.remove(f.name) - def list_dir_or_file(self, - dir_path: Union[str, Path], - list_dir: bool = True, - list_file: bool = True, - suffix: Optional[Union[str, Tuple[str]]] = None, - recursive: bool = False) -> Iterator[str]: + def list_dir_or_file( + self, + dir_path: Union[str, Path], + list_dir: bool = True, + list_file: bool = True, + suffix: Optional[Union[str, Tuple[str]]] = None, + recursive: bool = False, + ) -> Iterator[str]: """Scan a directory to find the interested directories or files in arbitrary order. @@ -360,54 +363,51 @@ def list_dir_or_file(self, Yields: Iterable[str]: A relative path to ``dir_path``. """ - if not has_method(self._client, 'list'): + if not has_method(self._client, "list"): raise NotImplementedError( - ('Current version of Petrel Python SDK has not supported ' - 'the `list` method, please use a higher version or dev' - ' branch instead.')) + ( + "Current version of Petrel Python SDK has not supported " + "the `list` method, please use a higher version or dev" + " branch instead." + ) + ) dir_path = self._map_path(dir_path) dir_path = self._format_path(dir_path) if list_dir and suffix is not None: - raise TypeError( - '`list_dir` should be False when `suffix` is not None') + raise TypeError("`list_dir` should be False when `suffix` is not None") if (suffix is not None) and not isinstance(suffix, (str, tuple)): - raise TypeError('`suffix` must be a string or tuple of strings') + raise TypeError("`suffix` must be a string or tuple of strings") # Petrel's simulated directory hierarchy assumes that directory paths # should end with `/` - if not dir_path.endswith('/'): - dir_path += '/' + if not dir_path.endswith("/"): + dir_path += "/" root = dir_path - def _list_dir_or_file(dir_path, list_dir, list_file, suffix, - recursive): + def _list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive): for path in self._client.list(dir_path): # the `self.isdir` is not used here to determine whether path # is a directory, because `self.isdir` relies on # `self._client.list` - if path.endswith('/'): # a directory path + if path.endswith("/"): # a directory path next_dir_path = self.join_path(dir_path, path) if list_dir: # get the relative path and exclude the last # character '/' - rel_dir = next_dir_path[len(root):-1] + rel_dir = next_dir_path[len(root) : -1] yield rel_dir if recursive: - yield from _list_dir_or_file(next_dir_path, list_dir, - list_file, suffix, - recursive) + yield from _list_dir_or_file(next_dir_path, list_dir, list_file, suffix, recursive) else: # a file path absolute_path = self.join_path(dir_path, path) - rel_path = absolute_path[len(root):] - if (suffix is None - or rel_path.endswith(suffix)) and list_file: + rel_path = absolute_path[len(root) :] + if (suffix is None or rel_path.endswith(suffix)) and list_file: yield rel_path - return _list_dir_or_file(dir_path, list_dir, list_file, suffix, - recursive) + return _list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive) class MemcachedBackend(BaseStorageBackend): @@ -423,23 +423,23 @@ class MemcachedBackend(BaseStorageBackend): def __init__(self, server_list_cfg, client_cfg, sys_path=None): if sys_path is not None: import sys + sys.path.append(sys_path) try: import mc except ImportError: - raise ImportError( - 'Please install memcached to enable MemcachedBackend.') + raise ImportError("Please install memcached to enable MemcachedBackend.") self.server_list_cfg = server_list_cfg self.client_cfg = client_cfg - self._client = mc.MemcachedClient.GetInstance(self.server_list_cfg, - self.client_cfg) + self._client = mc.MemcachedClient.GetInstance(self.server_list_cfg, self.client_cfg) # mc.pyvector servers as a point which points to a memory cache self._mc_buffer = mc.pyvector() def get(self, filepath): filepath = str(filepath) import mc + self._client.Get(filepath, self._mc_buffer) value_buf = mc.ConvertBuffer(self._mc_buffer) return value_buf @@ -466,24 +466,14 @@ class LmdbBackend(BaseStorageBackend): db_path (str): Lmdb database path. """ - def __init__(self, - db_path, - readonly=True, - lock=False, - readahead=False, - **kwargs): + def __init__(self, db_path, readonly=True, lock=False, readahead=False, **kwargs): try: import lmdb except ImportError: - raise ImportError('Please install lmdb to enable LmdbBackend.') + raise ImportError("Please install lmdb to enable LmdbBackend.") self.db_path = str(db_path) - self._client = lmdb.open( - self.db_path, - readonly=readonly, - lock=lock, - readahead=readahead, - **kwargs) + self._client = lmdb.open(self.db_path, readonly=readonly, lock=lock, readahead=readahead, **kwargs) def get(self, filepath): """Get values according to the filepath. @@ -493,7 +483,7 @@ def get(self, filepath): """ filepath = str(filepath) with self._client.begin(write=False) as txn: - value_buf = txn.get(filepath.encode('ascii')) + value_buf = txn.get(filepath.encode("ascii")) return value_buf def get_text(self, filepath, encoding=None): @@ -514,13 +504,11 @@ def get(self, filepath: Union[str, Path]) -> bytes: Returns: bytes: Expected bytes object. """ - with open(filepath, 'rb') as f: + with open(filepath, "rb") as f: value_buf = f.read() return value_buf - def get_text(self, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> str: + def get_text(self, filepath: Union[str, Path], encoding: str = "utf-8") -> str: """Read data from a given ``filepath`` with 'r' mode. Args: @@ -531,7 +519,7 @@ def get_text(self, Returns: str: Expected text reading from ``filepath``. """ - with open(filepath, 'r', encoding=encoding) as f: + with open(filepath, "r", encoding=encoding) as f: value_buf = f.read() return value_buf @@ -547,13 +535,10 @@ def put(self, obj: bytes, filepath: Union[str, Path]) -> None: filepath (str or Path): Path to write data. """ mmcv.mkdir_or_exist(osp.dirname(filepath)) - with open(filepath, 'wb') as f: + with open(filepath, "wb") as f: f.write(obj) - def put_text(self, - obj: str, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> None: + def put_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None: """Write data to a given ``filepath`` with 'w' mode. Note: @@ -567,7 +552,7 @@ def put_text(self, Default: 'utf-8'. """ mmcv.mkdir_or_exist(osp.dirname(filepath)) - with open(filepath, 'w', encoding=encoding) as f: + with open(filepath, "w", encoding=encoding) as f: f.write(obj) def remove(self, filepath: Union[str, Path]) -> None: @@ -614,8 +599,7 @@ def isfile(self, filepath: Union[str, Path]) -> bool: """ return osp.isfile(filepath) - def join_path(self, filepath: Union[str, Path], - *filepaths: Union[str, Path]) -> str: + def join_path(self, filepath: Union[str, Path], *filepaths: Union[str, Path]) -> str: """Concatenate all file paths. Join one or more filepath components intelligently. The return value @@ -630,17 +614,18 @@ def join_path(self, filepath: Union[str, Path], return osp.join(filepath, *filepaths) @contextmanager - def get_local_path( - self, filepath: Union[str, Path]) -> Iterable[Union[str, Path]]: + def get_local_path(self, filepath: Union[str, Path]) -> Iterable[Union[str, Path]]: """Only for unified API and do nothing.""" yield filepath - def list_dir_or_file(self, - dir_path: Union[str, Path], - list_dir: bool = True, - list_file: bool = True, - suffix: Optional[Union[str, Tuple[str]]] = None, - recursive: bool = False) -> Iterator[str]: + def list_dir_or_file( + self, + dir_path: Union[str, Path], + list_dir: bool = True, + list_file: bool = True, + suffix: Optional[Union[str, Tuple[str]]] = None, + recursive: bool = False, + ) -> Iterator[str]: """Scan a directory to find the interested directories or files in arbitrary order. @@ -660,32 +645,27 @@ def list_dir_or_file(self, Iterable[str]: A relative path to ``dir_path``. """ if list_dir and suffix is not None: - raise TypeError('`suffix` should be None when `list_dir` is True') + raise TypeError("`suffix` should be None when `list_dir` is True") if (suffix is not None) and not isinstance(suffix, (str, tuple)): - raise TypeError('`suffix` must be a string or tuple of strings') + raise TypeError("`suffix` must be a string or tuple of strings") root = dir_path - def _list_dir_or_file(dir_path, list_dir, list_file, suffix, - recursive): + def _list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive): for entry in os.scandir(dir_path): - if not entry.name.startswith('.') and entry.is_file(): + if not entry.name.startswith(".") and entry.is_file(): rel_path = osp.relpath(entry.path, root) - if (suffix is None - or rel_path.endswith(suffix)) and list_file: + if (suffix is None or rel_path.endswith(suffix)) and list_file: yield rel_path elif osp.isdir(entry.path): if list_dir: rel_dir = osp.relpath(entry.path, root) yield rel_dir if recursive: - yield from _list_dir_or_file(entry.path, list_dir, - list_file, suffix, - recursive) + yield from _list_dir_or_file(entry.path, list_dir, list_file, suffix, recursive) - return _list_dir_or_file(dir_path, list_dir, list_file, suffix, - recursive) + return _list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive) class HTTPBackend(BaseStorageBackend): @@ -695,7 +675,7 @@ def get(self, filepath): value_buf = urlopen(filepath).read() return value_buf - def get_text(self, filepath, encoding='utf-8'): + def get_text(self, filepath, encoding="utf-8"): value_buf = urlopen(filepath).read() return value_buf.decode(encoding) @@ -763,12 +743,12 @@ class FileClient: """ _backends = { - 'disk': HardDiskBackend, - 'ceph': CephBackend, - 'memcached': MemcachedBackend, - 'lmdb': LmdbBackend, - 'petrel': PetrelBackend, - 'http': HTTPBackend, + "disk": HardDiskBackend, + "ceph": CephBackend, + "memcached": MemcachedBackend, + "lmdb": LmdbBackend, + "petrel": PetrelBackend, + "http": HTTPBackend, } # This collection is used to record the overridden backends, and when a # backend appears in the collection, the singleton pattern is disabled for @@ -776,9 +756,9 @@ class FileClient: # returned will be the backend before overwriting _overridden_backends = set() _prefix_to_backends = { - 's3': PetrelBackend, - 'http': HTTPBackend, - 'https': HTTPBackend, + "s3": PetrelBackend, + "http": HTTPBackend, + "https": HTTPBackend, } _overridden_prefixes = set() @@ -786,26 +766,29 @@ class FileClient: def __new__(cls, backend=None, prefix=None, **kwargs): if backend is None and prefix is None: - backend = 'disk' + backend = "disk" if backend is not None and backend not in cls._backends: raise ValueError( - f'Backend {backend} is not supported. Currently supported ones' - f' are {list(cls._backends.keys())}') + f"Backend {backend} is not supported. Currently supported ones" f" are {list(cls._backends.keys())}" + ) if prefix is not None and prefix not in cls._prefix_to_backends: raise ValueError( - f'prefix {prefix} is not supported. Currently supported ones ' - f'are {list(cls._prefix_to_backends.keys())}') + f"prefix {prefix} is not supported. Currently supported ones " + f"are {list(cls._prefix_to_backends.keys())}" + ) # concatenate the arguments to a unique key for determining whether # objects with the same arguments were created - arg_key = f'{backend}:{prefix}' + arg_key = f"{backend}:{prefix}" for key, value in kwargs.items(): - arg_key += f':{key}:{value}' + arg_key += f":{key}:{value}" # if a backend was overridden, it will create a new object - if (arg_key in cls._instances - and backend not in cls._overridden_backends - and prefix not in cls._overridden_prefixes): + if ( + arg_key in cls._instances + and backend not in cls._overridden_backends + and prefix not in cls._overridden_prefixes + ): _instance = cls._instances[arg_key] else: # create a new object and put it to _instance @@ -844,20 +827,20 @@ def parse_uri_prefix(uri: Union[str, Path]) -> Optional[str]: """ assert is_filepath(uri) uri = str(uri) - if '://' not in uri: + if "://" not in uri: return None else: - prefix, _ = uri.split('://') + prefix, _ = uri.split("://") # In the case of PetrelBackend, the prefix may contains the cluster # name like clusterName:s3 - if ':' in prefix: - _, prefix = prefix.split(':') + if ":" in prefix: + _, prefix = prefix.split(":") return prefix @classmethod - def infer_client(cls, - file_client_args: Optional[dict] = None, - uri: Optional[Union[str, Path]] = None) -> 'FileClient': + def infer_client( + cls, file_client_args: Optional[dict] = None, uri: Optional[Union[str, Path]] = None + ) -> "FileClient": """Infer a suitable file client based on the URI and arguments. Args: @@ -885,18 +868,15 @@ def infer_client(cls, @classmethod def _register_backend(cls, name, backend, force=False, prefixes=None): if not isinstance(name, str): - raise TypeError('the backend name should be a string, ' - f'but got {type(name)}') + raise TypeError("the backend name should be a string, " f"but got {type(name)}") if not inspect.isclass(backend): - raise TypeError( - f'backend should be a class but got {type(backend)}') + raise TypeError(f"backend should be a class but got {type(backend)}") if not issubclass(backend, BaseStorageBackend): - raise TypeError( - f'backend {backend} is not a subclass of BaseStorageBackend') + raise TypeError(f"backend {backend} is not a subclass of BaseStorageBackend") if not force and name in cls._backends: raise KeyError( - f'{name} is already registered as a storage backend, ' - 'add "force=True" if you want to override it') + f"{name} is already registered as a storage backend, " 'add "force=True" if you want to override it' + ) if name in cls._backends and force: cls._overridden_backends.add(name) @@ -915,8 +895,9 @@ def _register_backend(cls, name, backend, force=False, prefixes=None): cls._prefix_to_backends[prefix] = backend else: raise KeyError( - f'{prefix} is already registered as a storage backend,' - ' add "force=True" if you want to override it') + f"{prefix} is already registered as a storage backend," + ' add "force=True" if you want to override it' + ) @classmethod def register_backend(cls, name, backend=None, force=False, prefixes=None): @@ -962,13 +943,11 @@ def get_text(self, filepath): `New in version 1.3.15.` """ if backend is not None: - cls._register_backend( - name, backend, force=force, prefixes=prefixes) + cls._register_backend(name, backend, force=force, prefixes=prefixes) return def _register(backend_cls): - cls._register_backend( - name, backend_cls, force=force, prefixes=prefixes) + cls._register_backend(name, backend_cls, force=force, prefixes=prefixes) return backend_cls return _register @@ -991,7 +970,7 @@ def get(self, filepath: Union[str, Path]) -> Union[bytes, memoryview]: """ return self.client.get(filepath) - def get_text(self, filepath: Union[str, Path], encoding='utf-8') -> str: + def get_text(self, filepath: Union[str, Path], encoding="utf-8") -> str: """Read data from a given ``filepath`` with 'r' mode. Args: @@ -1076,8 +1055,7 @@ def isfile(self, filepath: Union[str, Path]) -> bool: """ return self.client.isfile(filepath) - def join_path(self, filepath: Union[str, Path], - *filepaths: Union[str, Path]) -> str: + def join_path(self, filepath: Union[str, Path], *filepaths: Union[str, Path]) -> str: """Concatenate all file paths. Join one or more filepath components intelligently. The return value @@ -1120,12 +1098,14 @@ def get_local_path(self, filepath: Union[str, Path]) -> Iterable[str]: with self.client.get_local_path(str(filepath)) as local_path: yield local_path - def list_dir_or_file(self, - dir_path: Union[str, Path], - list_dir: bool = True, - list_file: bool = True, - suffix: Optional[Union[str, Tuple[str]]] = None, - recursive: bool = False) -> Iterator[str]: + def list_dir_or_file( + self, + dir_path: Union[str, Path], + list_dir: bool = True, + list_file: bool = True, + suffix: Optional[Union[str, Tuple[str]]] = None, + recursive: bool = False, + ) -> Iterator[str]: """Scan a directory to find the interested directories or files in arbitrary order. @@ -1144,5 +1124,4 @@ def list_dir_or_file(self, Yields: Iterable[str]: A relative path to ``dir_path``. """ - yield from self.client.list_dir_or_file(dir_path, list_dir, list_file, - suffix, recursive) + yield from self.client.list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/image/io.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/image/io.py index b06ad407..9e745a31 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/image/io.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/image/io.py @@ -5,10 +5,9 @@ import cv2 import numpy as np -from cv2 import (IMREAD_COLOR, IMREAD_GRAYSCALE, IMREAD_IGNORE_ORIENTATION, - IMREAD_UNCHANGED) +from cv2 import IMREAD_COLOR, IMREAD_GRAYSCALE, IMREAD_IGNORE_ORIENTATION, IMREAD_UNCHANGED -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import check_file_exist, is_str, mkdir_or_exist +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import check_file_exist, is_str, mkdir_or_exist try: from turbojpeg import TJCS_RGB, TJPF_BGR, TJPF_GRAY, TurboJPEG @@ -26,18 +25,17 @@ tifffile = None jpeg = None -supported_backends = ['cv2', 'turbojpeg', 'pillow', 'tifffile'] +supported_backends = ["cv2", "turbojpeg", "pillow", "tifffile"] imread_flags = { - 'color': IMREAD_COLOR, - 'grayscale': IMREAD_GRAYSCALE, - 'unchanged': IMREAD_UNCHANGED, - 'color_ignore_orientation': IMREAD_IGNORE_ORIENTATION | IMREAD_COLOR, - 'grayscale_ignore_orientation': - IMREAD_IGNORE_ORIENTATION | IMREAD_GRAYSCALE + "color": IMREAD_COLOR, + "grayscale": IMREAD_GRAYSCALE, + "unchanged": IMREAD_UNCHANGED, + "color_ignore_orientation": IMREAD_IGNORE_ORIENTATION | IMREAD_COLOR, + "grayscale_ignore_orientation": IMREAD_IGNORE_ORIENTATION | IMREAD_GRAYSCALE, } -imread_backend = 'cv2' +imread_backend = "cv2" def use_backend(backend): @@ -52,37 +50,37 @@ def use_backend(backend): assert backend in supported_backends global imread_backend imread_backend = backend - if imread_backend == 'turbojpeg': + if imread_backend == "turbojpeg": if TurboJPEG is None: - raise ImportError('`PyTurboJPEG` is not installed') + raise ImportError("`PyTurboJPEG` is not installed") global jpeg if jpeg is None: jpeg = TurboJPEG() - elif imread_backend == 'pillow': + elif imread_backend == "pillow": if Image is None: - raise ImportError('`Pillow` is not installed') - elif imread_backend == 'tifffile': + raise ImportError("`Pillow` is not installed") + elif imread_backend == "tifffile": if tifffile is None: - raise ImportError('`tifffile` is not installed') + raise ImportError("`tifffile` is not installed") -def _jpegflag(flag='color', channel_order='bgr'): +def _jpegflag(flag="color", channel_order="bgr"): channel_order = channel_order.lower() - if channel_order not in ['rgb', 'bgr']: + if channel_order not in ["rgb", "bgr"]: raise ValueError('channel order must be either "rgb" or "bgr"') - if flag == 'color': - if channel_order == 'bgr': + if flag == "color": + if channel_order == "bgr": return TJPF_BGR - elif channel_order == 'rgb': + elif channel_order == "rgb": return TJCS_RGB - elif flag == 'grayscale': + elif flag == "grayscale": return TJPF_GRAY else: raise ValueError('flag must be "color" or "grayscale"') -def _pillow2array(img, flag='color', channel_order='bgr'): +def _pillow2array(img, flag="color", channel_order="bgr"): """Convert a pillow image to numpy array. Args: @@ -97,47 +95,48 @@ def _pillow2array(img, flag='color', channel_order='bgr'): np.ndarray: The converted numpy array """ channel_order = channel_order.lower() - if channel_order not in ['rgb', 'bgr']: + if channel_order not in ["rgb", "bgr"]: raise ValueError('channel order must be either "rgb" or "bgr"') - if flag == 'unchanged': + if flag == "unchanged": array = np.array(img) if array.ndim >= 3 and array.shape[2] >= 3: # color image array[:, :, :3] = array[:, :, (2, 1, 0)] # RGB to BGR else: # Handle exif orientation tag - if flag in ['color', 'grayscale']: + if flag in ["color", "grayscale"]: img = ImageOps.exif_transpose(img) # If the image mode is not 'RGB', convert it to 'RGB' first. - if img.mode != 'RGB': - if img.mode != 'LA': + if img.mode != "RGB": + if img.mode != "LA": # Most formats except 'LA' can be directly converted to RGB - img = img.convert('RGB') + img = img.convert("RGB") else: # When the mode is 'LA', the default conversion will fill in # the canvas with black, which sometimes shadows black objects # in the foreground. # # Therefore, a random color (124, 117, 104) is used for canvas - img_rgba = img.convert('RGBA') - img = Image.new('RGB', img_rgba.size, (124, 117, 104)) + img_rgba = img.convert("RGBA") + img = Image.new("RGB", img_rgba.size, (124, 117, 104)) img.paste(img_rgba, mask=img_rgba.split()[3]) # 3 is alpha - if flag in ['color', 'color_ignore_orientation']: + if flag in ["color", "color_ignore_orientation"]: array = np.array(img) - if channel_order != 'rgb': + if channel_order != "rgb": array = array[:, :, ::-1] # RGB to BGR - elif flag in ['grayscale', 'grayscale_ignore_orientation']: - img = img.convert('L') + elif flag in ["grayscale", "grayscale_ignore_orientation"]: + img = img.convert("L") array = np.array(img) else: raise ValueError( 'flag must be "color", "grayscale", "unchanged", ' f'"color_ignore_orientation" or "grayscale_ignore_orientation"' - f' but got {flag}') + f" but got {flag}" + ) return array -def imread(img_or_path, flag='color', channel_order='bgr', backend=None): +def imread(img_or_path, flag="color", channel_order="bgr", backend=None): """Read an image. Args: @@ -165,42 +164,40 @@ def imread(img_or_path, flag='color', channel_order='bgr', backend=None): if backend is None: backend = imread_backend if backend not in supported_backends: - raise ValueError(f'backend: {backend} is not supported. Supported ' - "backends are 'cv2', 'turbojpeg', 'pillow'") + raise ValueError( + f"backend: {backend} is not supported. Supported " "backends are 'cv2', 'turbojpeg', 'pillow'" + ) if isinstance(img_or_path, Path): img_or_path = str(img_or_path) if isinstance(img_or_path, np.ndarray): return img_or_path elif is_str(img_or_path): - check_file_exist(img_or_path, - f'img file does not exist: {img_or_path}') - if backend == 'turbojpeg': - with open(img_or_path, 'rb') as in_file: - img = jpeg.decode(in_file.read(), - _jpegflag(flag, channel_order)) + check_file_exist(img_or_path, f"img file does not exist: {img_or_path}") + if backend == "turbojpeg": + with open(img_or_path, "rb") as in_file: + img = jpeg.decode(in_file.read(), _jpegflag(flag, channel_order)) if img.shape[-1] == 1: img = img[:, :, 0] return img - elif backend == 'pillow': + elif backend == "pillow": img = Image.open(img_or_path) img = _pillow2array(img, flag, channel_order) return img - elif backend == 'tifffile': + elif backend == "tifffile": img = tifffile.imread(img_or_path) return img else: flag = imread_flags[flag] if is_str(flag) else flag img = cv2.imread(img_or_path, flag) - if flag == IMREAD_COLOR and channel_order == 'rgb': + if flag == IMREAD_COLOR and channel_order == "rgb": cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img) return img else: - raise TypeError('"img" must be a numpy array or a str or ' - 'a pathlib.Path object') + raise TypeError('"img" must be a numpy array or a str or ' "a pathlib.Path object") -def imfrombytes(content, flag='color', channel_order='bgr', backend=None): +def imfrombytes(content, flag="color", channel_order="bgr", backend=None): """Read an image from bytes. Args: @@ -218,14 +215,15 @@ def imfrombytes(content, flag='color', channel_order='bgr', backend=None): if backend is None: backend = imread_backend if backend not in supported_backends: - raise ValueError(f'backend: {backend} is not supported. Supported ' - "backends are 'cv2', 'turbojpeg', 'pillow'") - if backend == 'turbojpeg': + raise ValueError( + f"backend: {backend} is not supported. Supported " "backends are 'cv2', 'turbojpeg', 'pillow'" + ) + if backend == "turbojpeg": img = jpeg.decode(content, _jpegflag(flag, channel_order)) if img.shape[-1] == 1: img = img[:, :, 0] return img - elif backend == 'pillow': + elif backend == "pillow": buff = io.BytesIO(content) img = Image.open(buff) img = _pillow2array(img, flag, channel_order) @@ -234,7 +232,7 @@ def imfrombytes(content, flag='color', channel_order='bgr', backend=None): img_np = np.frombuffer(content, np.uint8) flag = imread_flags[flag] if is_str(flag) else flag img = cv2.imdecode(img_np, flag) - if flag == IMREAD_COLOR and channel_order == 'rgb': + if flag == IMREAD_COLOR and channel_order == "rgb": cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img) return img diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/image/misc.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/image/misc.py index 845ed10b..a9b9e3c7 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/image/misc.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/image/misc.py @@ -1,7 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. import numpy as np -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv try: import torch @@ -27,7 +27,7 @@ def tensor2imgs(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True): """ if torch is None: - raise RuntimeError('pytorch is not installed') + raise RuntimeError("pytorch is not installed") assert torch.is_tensor(tensor) and tensor.ndim == 4 assert len(mean) == 3 assert len(std) == 3 @@ -38,7 +38,6 @@ def tensor2imgs(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True): imgs = [] for img_id in range(num_imgs): img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0) - img = mmcv.imdenormalize( - img, mean, std, to_bgr=to_rgb).astype(np.uint8) + img = mmcv.imdenormalize(img, mean, std, to_bgr=to_rgb).astype(np.uint8) imgs.append(np.ascontiguousarray(img)) return imgs diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/cc_attention.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/cc_attention.py index 9c1727ed..3d0f226e 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/cc_attention.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/cc_attention.py @@ -3,7 +3,7 @@ import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import PLUGIN_LAYERS, Scale +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import PLUGIN_LAYERS, Scale def NEG_INF_DIAG(n, device): @@ -12,7 +12,7 @@ def NEG_INF_DIAG(n, device): The diagonal are all "-inf". This is for avoiding calculating the overlapped element in the Criss-Cross twice. """ - return torch.diag(torch.tensor(float('-inf')).to(device).repeat(n), 0) + return torch.diag(torch.tensor(float("-inf")).to(device).repeat(n), 0) @PLUGIN_LAYERS.register_module() @@ -46,7 +46,7 @@ def __init__(self, in_channels): self.query_conv = nn.Conv2d(in_channels, in_channels // 8, 1) self.key_conv = nn.Conv2d(in_channels, in_channels // 8, 1) self.value_conv = nn.Conv2d(in_channels, in_channels, 1) - self.gamma = Scale(0.) + self.gamma = Scale(0.0) self.in_channels = in_channels def forward(self, x): @@ -63,14 +63,12 @@ def forward(self, x): query = self.query_conv(x) key = self.key_conv(x) value = self.value_conv(x) - energy_H = torch.einsum('bchw,bciw->bwhi', query, key) + NEG_INF_DIAG( - H, query.device) + energy_H = torch.einsum("bchw,bciw->bwhi", query, key) + NEG_INF_DIAG(H, query.device) energy_H = energy_H.transpose(1, 2) - energy_W = torch.einsum('bchw,bchj->bhwj', query, key) - attn = F.softmax( - torch.cat([energy_H, energy_W], dim=-1), dim=-1) # [B,H,W,(H+W)] - out = torch.einsum('bciw,bhwi->bchw', value, attn[..., :H]) - out += torch.einsum('bchj,bhwj->bchw', value, attn[..., H:]) + energy_W = torch.einsum("bchw,bchj->bhwj", query, key) + attn = F.softmax(torch.cat([energy_H, energy_W], dim=-1), dim=-1) # [B,H,W,(H+W)] + out = torch.einsum("bciw,bhwi->bchw", value, attn[..., :H]) + out += torch.einsum("bchj,bhwj->bchw", value, attn[..., H:]) out = self.gamma(out) + x out = out.contiguous() @@ -79,5 +77,5 @@ def forward(self, x): def __repr__(self): s = self.__class__.__name__ - s += f'(in_channels={self.in_channels})' + s += f"(in_channels={self.in_channels})" return s diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/deform_conv.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/deform_conv.py index 9bb7abaa..baf71185 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/deform_conv.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/deform_conv.py @@ -9,32 +9,22 @@ from torch.autograd.function import once_differentiable from torch.nn.modules.utils import _pair, _single -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import deprecated_api_warning +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import deprecated_api_warning from ..cnn import CONV_LAYERS from ..utils import ext_loader, print_log -ext_module = ext_loader.load_ext('_ext', [ - 'deform_conv_forward', 'deform_conv_backward_input', - 'deform_conv_backward_parameters' -]) +ext_module = ext_loader.load_ext( + "_ext", ["deform_conv_forward", "deform_conv_backward_input", "deform_conv_backward_parameters"] +) class DeformConv2dFunction(Function): - @staticmethod - def symbolic(g, - input, - offset, - weight, - stride, - padding, - dilation, - groups, - deform_groups, - bias=False, - im2col_step=32): + def symbolic( + g, input, offset, weight, stride, padding, dilation, groups, deform_groups, bias=False, im2col_step=32 + ): return g.op( - 'mmcv::MMCVDeformConv2d', + "mmcv::MMCVDeformConv2d", input, offset, weight, @@ -44,25 +34,29 @@ def symbolic(g, groups_i=groups, deform_groups_i=deform_groups, bias_i=bias, - im2col_step_i=im2col_step) + im2col_step_i=im2col_step, + ) @staticmethod - def forward(ctx, - input, - offset, - weight, - stride=1, - padding=0, - dilation=1, - groups=1, - deform_groups=1, - bias=False, - im2col_step=32): + def forward( + ctx, + input, + offset, + weight, + stride=1, + padding=0, + dilation=1, + groups=1, + deform_groups=1, + bias=False, + im2col_step=32, + ): if input is not None and input.dim() != 4: raise ValueError( - f'Expected 4D tensor as input, got {input.dim()}D tensor \ - instead.') - assert bias is False, 'Only support bias is False.' + f"Expected 4D tensor as input, got {input.dim()}D tensor \ + instead." + ) + assert bias is False, "Only support bias is False." ctx.stride = _pair(stride) ctx.padding = _pair(padding) ctx.dilation = _pair(dilation) @@ -81,14 +75,12 @@ def forward(ctx, weight = weight.type_as(input) ctx.save_for_backward(input, offset, weight) - output = input.new_empty( - DeformConv2dFunction._output_size(ctx, input, weight)) + output = input.new_empty(DeformConv2dFunction._output_size(ctx, input, weight)) ctx.bufs_ = [input.new_empty(0), input.new_empty(0)] # columns, ones cur_im2col_step = min(ctx.im2col_step, input.size(0)) - assert (input.size(0) % - cur_im2col_step) == 0, 'im2col step must divide batchsize' + assert (input.size(0) % cur_im2col_step) == 0, "im2col step must divide batchsize" ext_module.deform_conv_forward( input, weight, @@ -106,7 +98,8 @@ def forward(ctx, dilationH=ctx.dilation[0], group=ctx.groups, deformable_group=ctx.deform_groups, - im2col_step=cur_im2col_step) + im2col_step=cur_im2col_step, + ) return output @staticmethod @@ -117,8 +110,7 @@ def backward(ctx, grad_output): grad_input = grad_offset = grad_weight = None cur_im2col_step = min(ctx.im2col_step, input.size(0)) - assert (input.size(0) % cur_im2col_step - ) == 0, 'batch size must be divisible by im2col_step' + assert (input.size(0) % cur_im2col_step) == 0, "batch size must be divisible by im2col_step" grad_output = grad_output.contiguous() if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]: @@ -142,7 +134,8 @@ def backward(ctx, grad_output): dilationH=ctx.dilation[0], group=ctx.groups, deformable_group=ctx.deform_groups, - im2col_step=cur_im2col_step) + im2col_step=cur_im2col_step, + ) if ctx.needs_input_grad[2]: grad_weight = torch.zeros_like(weight) @@ -164,10 +157,10 @@ def backward(ctx, grad_output): group=ctx.groups, deformable_group=ctx.deform_groups, scale=1, - im2col_step=cur_im2col_step) + im2col_step=cur_im2col_step, + ) - return grad_input, grad_offset, grad_weight, \ - None, None, None, None, None, None, None + return grad_input, grad_offset, grad_weight, None, None, None, None, None, None, None @staticmethod def _output_size(ctx, input, weight): @@ -178,11 +171,11 @@ def _output_size(ctx, input, weight): pad = ctx.padding[d] kernel = ctx.dilation[d] * (weight.size(d + 2) - 1) + 1 stride_ = ctx.stride[d] - output_size += ((in_size + (2 * pad) - kernel) // stride_ + 1, ) + output_size += ((in_size + (2 * pad) - kernel) // stride_ + 1,) if not all(map(lambda s: s > 0, output_size)): raise ValueError( - 'convolution input is too small (output would be ' + - 'x'.join(map(str, output_size)) + ')') + "convolution input is too small (output would be " + "x".join(map(str, output_size)) + ")" + ) return output_size @@ -223,28 +216,28 @@ class DeformConv2d(nn.Module): `New in version 1.3.17.` """ - @deprecated_api_warning({'deformable_groups': 'deform_groups'}, - cls_name='DeformConv2d') - def __init__(self, - in_channels: int, - out_channels: int, - kernel_size: Union[int, Tuple[int, ...]], - stride: Union[int, Tuple[int, ...]] = 1, - padding: Union[int, Tuple[int, ...]] = 0, - dilation: Union[int, Tuple[int, ...]] = 1, - groups: int = 1, - deform_groups: int = 1, - bias: bool = False, - im2col_step: int = 32) -> None: + @deprecated_api_warning({"deformable_groups": "deform_groups"}, cls_name="DeformConv2d") + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, ...]], + stride: Union[int, Tuple[int, ...]] = 1, + padding: Union[int, Tuple[int, ...]] = 0, + dilation: Union[int, Tuple[int, ...]] = 1, + groups: int = 1, + deform_groups: int = 1, + bias: bool = False, + im2col_step: int = 32, + ) -> None: super(DeformConv2d, self).__init__() - assert not bias, \ - f'bias={bias} is not supported in DeformConv2d.' - assert in_channels % groups == 0, \ - f'in_channels {in_channels} cannot be divisible by groups {groups}' - assert out_channels % groups == 0, \ - f'out_channels {out_channels} cannot be divisible by groups \ - {groups}' + assert not bias, f"bias={bias} is not supported in DeformConv2d." + assert in_channels % groups == 0, f"in_channels {in_channels} cannot be divisible by groups {groups}" + assert ( + out_channels % groups == 0 + ), f"out_channels {out_channels} cannot be divisible by groups \ + {groups}" self.in_channels = in_channels self.out_channels = out_channels @@ -260,9 +253,7 @@ def __init__(self, self.output_padding = _single(0) # only weight, no bias - self.weight = nn.Parameter( - torch.Tensor(out_channels, in_channels // self.groups, - *self.kernel_size)) + self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels // self.groups, *self.kernel_size)) self.reset_parameters() @@ -271,7 +262,7 @@ def reset_parameters(self): # method described in `Delving deep into rectifiers: Surpassing # human-level performance on ImageNet classification` - He, K. et al. # (2015), using a uniform distribution - nn.init.kaiming_uniform_(self.weight, nonlinearity='relu') + nn.init.kaiming_uniform_(self.weight, nonlinearity="relu") def forward(self, x: Tensor, offset: Tensor) -> Tensor: """Deformable Convolutional forward function. @@ -296,38 +287,45 @@ def forward(self, x: Tensor, offset: Tensor) -> Tensor: """ # To fix an assert error in deform_conv_cuda.cpp:128 # input image is smaller than kernel - input_pad = (x.size(2) < self.kernel_size[0]) or (x.size(3) < - self.kernel_size[1]) + input_pad = (x.size(2) < self.kernel_size[0]) or (x.size(3) < self.kernel_size[1]) if input_pad: pad_h = max(self.kernel_size[0] - x.size(2), 0) pad_w = max(self.kernel_size[1] - x.size(3), 0) - x = F.pad(x, (0, pad_w, 0, pad_h), 'constant', 0).contiguous() - offset = F.pad(offset, (0, pad_w, 0, pad_h), 'constant', 0) + x = F.pad(x, (0, pad_w, 0, pad_h), "constant", 0).contiguous() + offset = F.pad(offset, (0, pad_w, 0, pad_h), "constant", 0) offset = offset.contiguous() - out = deform_conv2d(x, offset, self.weight, self.stride, self.padding, - self.dilation, self.groups, self.deform_groups, - False, self.im2col_step) + out = deform_conv2d( + x, + offset, + self.weight, + self.stride, + self.padding, + self.dilation, + self.groups, + self.deform_groups, + False, + self.im2col_step, + ) if input_pad: - out = out[:, :, :out.size(2) - pad_h, :out.size(3) - - pad_w].contiguous() + out = out[:, :, : out.size(2) - pad_h, : out.size(3) - pad_w].contiguous() return out def __repr__(self): s = self.__class__.__name__ - s += f'(in_channels={self.in_channels},\n' - s += f'out_channels={self.out_channels},\n' - s += f'kernel_size={self.kernel_size},\n' - s += f'stride={self.stride},\n' - s += f'padding={self.padding},\n' - s += f'dilation={self.dilation},\n' - s += f'groups={self.groups},\n' - s += f'deform_groups={self.deform_groups},\n' + s += f"(in_channels={self.in_channels},\n" + s += f"out_channels={self.out_channels},\n" + s += f"kernel_size={self.kernel_size},\n" + s += f"stride={self.stride},\n" + s += f"padding={self.padding},\n" + s += f"dilation={self.dilation},\n" + s += f"groups={self.groups},\n" + s += f"deform_groups={self.deform_groups},\n" # bias is not supported in DeformConv2d. - s += 'bias=False)' + s += "bias=False)" return s -@CONV_LAYERS.register_module('DCN') +@CONV_LAYERS.register_module("DCN") class DeformConv2dPack(DeformConv2d): """A Deformable Conv Encapsulation that acts as normal Conv layers. @@ -364,7 +362,8 @@ def __init__(self, *args, **kwargs): stride=_pair(self.stride), padding=_pair(self.padding), dilation=_pair(self.dilation), - bias=True) + bias=True, + ) self.init_offset() def init_offset(self): @@ -373,33 +372,35 @@ def init_offset(self): def forward(self, x): offset = self.conv_offset(x) - return deform_conv2d(x, offset, self.weight, self.stride, self.padding, - self.dilation, self.groups, self.deform_groups, - False, self.im2col_step) - - def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, - missing_keys, unexpected_keys, error_msgs): - version = local_metadata.get('version', None) + return deform_conv2d( + x, + offset, + self.weight, + self.stride, + self.padding, + self.dilation, + self.groups, + self.deform_groups, + False, + self.im2col_step, + ) + + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + version = local_metadata.get("version", None) if version is None or version < 2: # the key is different in early versions # In version < 2, DeformConvPack loads previous benchmark models. - if (prefix + 'conv_offset.weight' not in state_dict - and prefix[:-1] + '_offset.weight' in state_dict): - state_dict[prefix + 'conv_offset.weight'] = state_dict.pop( - prefix[:-1] + '_offset.weight') - if (prefix + 'conv_offset.bias' not in state_dict - and prefix[:-1] + '_offset.bias' in state_dict): - state_dict[prefix + - 'conv_offset.bias'] = state_dict.pop(prefix[:-1] + - '_offset.bias') + if prefix + "conv_offset.weight" not in state_dict and prefix[:-1] + "_offset.weight" in state_dict: + state_dict[prefix + "conv_offset.weight"] = state_dict.pop(prefix[:-1] + "_offset.weight") + if prefix + "conv_offset.bias" not in state_dict and prefix[:-1] + "_offset.bias" in state_dict: + state_dict[prefix + "conv_offset.bias"] = state_dict.pop(prefix[:-1] + "_offset.bias") if version is not None and version > 1: - print_log( - f'DeformConv2dPack {prefix.rstrip(".")} is upgraded to ' - 'version 2.', - logger='root') - - super()._load_from_state_dict(state_dict, prefix, local_metadata, - strict, missing_keys, unexpected_keys, - error_msgs) + print_log(f'DeformConv2dPack {prefix.rstrip(".")} is upgraded to ' "version 2.", logger="root") + + super()._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/modulated_deform_conv.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/modulated_deform_conv.py index 9515ae91..0b899a0d 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/modulated_deform_conv.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/modulated_deform_conv.py @@ -7,48 +7,38 @@ from torch.autograd.function import once_differentiable from torch.nn.modules.utils import _pair, _single -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import deprecated_api_warning +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import deprecated_api_warning from ..cnn import CONV_LAYERS from ..utils import ext_loader, print_log -ext_module = ext_loader.load_ext( - '_ext', - ['modulated_deform_conv_forward', 'modulated_deform_conv_backward']) +ext_module = ext_loader.load_ext("_ext", ["modulated_deform_conv_forward", "modulated_deform_conv_backward"]) class ModulatedDeformConv2dFunction(Function): - @staticmethod - def symbolic(g, input, offset, mask, weight, bias, stride, padding, - dilation, groups, deform_groups): + def symbolic(g, input, offset, mask, weight, bias, stride, padding, dilation, groups, deform_groups): input_tensors = [input, offset, mask, weight] if bias is not None: input_tensors.append(bias) return g.op( - 'mmcv::MMCVModulatedDeformConv2d', + "mmcv::MMCVModulatedDeformConv2d", *input_tensors, stride_i=stride, padding_i=padding, dilation_i=dilation, groups_i=groups, - deform_groups_i=deform_groups) + deform_groups_i=deform_groups, + ) @staticmethod - def forward(ctx, - input, - offset, - mask, - weight, - bias=None, - stride=1, - padding=0, - dilation=1, - groups=1, - deform_groups=1): + def forward( + ctx, input, offset, mask, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, deform_groups=1 + ): if input is not None and input.dim() != 4: raise ValueError( - f'Expected 4D tensor as input, got {input.dim()}D tensor \ - instead.') + f"Expected 4D tensor as input, got {input.dim()}D tensor \ + instead." + ) ctx.stride = _pair(stride) ctx.padding = _pair(padding) ctx.dilation = _pair(dilation) @@ -67,8 +57,7 @@ def forward(ctx, input = input.type_as(offset) weight = weight.type_as(input) ctx.save_for_backward(input, offset, mask, weight, bias) - output = input.new_empty( - ModulatedDeformConv2dFunction._output_size(ctx, input, weight)) + output = input.new_empty(ModulatedDeformConv2dFunction._output_size(ctx, input, weight)) ctx._bufs = [input.new_empty(0), input.new_empty(0)] ext_module.modulated_deform_conv_forward( input, @@ -89,7 +78,8 @@ def forward(ctx, dilation_w=ctx.dilation[1], group=ctx.groups, deformable_group=ctx.deform_groups, - with_bias=ctx.with_bias) + with_bias=ctx.with_bias, + ) return output @staticmethod @@ -126,12 +116,12 @@ def backward(ctx, grad_output): dilation_w=ctx.dilation[1], group=ctx.groups, deformable_group=ctx.deform_groups, - with_bias=ctx.with_bias) + with_bias=ctx.with_bias, + ) if not ctx.with_bias: grad_bias = None - return (grad_input, grad_offset, grad_mask, grad_weight, grad_bias, - None, None, None, None, None) + return (grad_input, grad_offset, grad_mask, grad_weight, grad_bias, None, None, None, None, None) @staticmethod def _output_size(ctx, input, weight): @@ -142,11 +132,11 @@ def _output_size(ctx, input, weight): pad = ctx.padding[d] kernel = ctx.dilation[d] * (weight.size(d + 2) - 1) + 1 stride_ = ctx.stride[d] - output_size += ((in_size + (2 * pad) - kernel) // stride_ + 1, ) + output_size += ((in_size + (2 * pad) - kernel) // stride_ + 1,) if not all(map(lambda s: s > 0, output_size)): raise ValueError( - 'convolution input is too small (output would be ' + - 'x'.join(map(str, output_size)) + ')') + "convolution input is too small (output would be " + "x".join(map(str, output_size)) + ")" + ) return output_size @@ -154,19 +144,19 @@ def _output_size(ctx, input, weight): class ModulatedDeformConv2d(nn.Module): - - @deprecated_api_warning({'deformable_groups': 'deform_groups'}, - cls_name='ModulatedDeformConv2d') - def __init__(self, - in_channels, - out_channels, - kernel_size, - stride=1, - padding=0, - dilation=1, - groups=1, - deform_groups=1, - bias=True): + @deprecated_api_warning({"deformable_groups": "deform_groups"}, cls_name="ModulatedDeformConv2d") + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + groups=1, + deform_groups=1, + bias=True, + ): super(ModulatedDeformConv2d, self).__init__() self.in_channels = in_channels self.out_channels = out_channels @@ -180,32 +170,38 @@ def __init__(self, self.transposed = False self.output_padding = _single(0) - self.weight = nn.Parameter( - torch.Tensor(out_channels, in_channels // groups, - *self.kernel_size)) + self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels // groups, *self.kernel_size)) if bias: self.bias = nn.Parameter(torch.Tensor(out_channels)) else: - self.register_parameter('bias', None) + self.register_parameter("bias", None) self.init_weights() def init_weights(self): n = self.in_channels for k in self.kernel_size: n *= k - stdv = 1. / math.sqrt(n) + stdv = 1.0 / math.sqrt(n) self.weight.data.uniform_(-stdv, stdv) if self.bias is not None: self.bias.data.zero_() def forward(self, x, offset, mask): - return modulated_deform_conv2d(x, offset, mask, self.weight, self.bias, - self.stride, self.padding, - self.dilation, self.groups, - self.deform_groups) + return modulated_deform_conv2d( + x, + offset, + mask, + self.weight, + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, + self.deform_groups, + ) -@CONV_LAYERS.register_module('DCNv2') +@CONV_LAYERS.register_module("DCNv2") class ModulatedDeformConv2dPack(ModulatedDeformConv2d): """A ModulatedDeformable Conv Encapsulation that acts as normal Conv layers. @@ -234,12 +230,13 @@ def __init__(self, *args, **kwargs): stride=self.stride, padding=self.padding, dilation=self.dilation, - bias=True) + bias=True, + ) self.init_weights() def init_weights(self): super(ModulatedDeformConv2dPack, self).init_weights() - if hasattr(self, 'conv_offset'): + if hasattr(self, "conv_offset"): self.conv_offset.weight.data.zero_() self.conv_offset.bias.data.zero_() @@ -248,35 +245,36 @@ def forward(self, x): o1, o2, mask = torch.chunk(out, 3, dim=1) offset = torch.cat((o1, o2), dim=1) mask = torch.sigmoid(mask) - return modulated_deform_conv2d(x, offset, mask, self.weight, self.bias, - self.stride, self.padding, - self.dilation, self.groups, - self.deform_groups) + return modulated_deform_conv2d( + x, + offset, + mask, + self.weight, + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, + self.deform_groups, + ) - def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, - missing_keys, unexpected_keys, error_msgs): - version = local_metadata.get('version', None) + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + version = local_metadata.get("version", None) if version is None or version < 2: # the key is different in early versions # In version < 2, ModulatedDeformConvPack # loads previous benchmark models. - if (prefix + 'conv_offset.weight' not in state_dict - and prefix[:-1] + '_offset.weight' in state_dict): - state_dict[prefix + 'conv_offset.weight'] = state_dict.pop( - prefix[:-1] + '_offset.weight') - if (prefix + 'conv_offset.bias' not in state_dict - and prefix[:-1] + '_offset.bias' in state_dict): - state_dict[prefix + - 'conv_offset.bias'] = state_dict.pop(prefix[:-1] + - '_offset.bias') + if prefix + "conv_offset.weight" not in state_dict and prefix[:-1] + "_offset.weight" in state_dict: + state_dict[prefix + "conv_offset.weight"] = state_dict.pop(prefix[:-1] + "_offset.weight") + if prefix + "conv_offset.bias" not in state_dict and prefix[:-1] + "_offset.bias" in state_dict: + state_dict[prefix + "conv_offset.bias"] = state_dict.pop(prefix[:-1] + "_offset.bias") if version is not None and version > 1: - print_log( - f'ModulatedDeformConvPack {prefix.rstrip(".")} is upgraded to ' - 'version 2.', - logger='root') + print_log(f'ModulatedDeformConvPack {prefix.rstrip(".")} is upgraded to ' "version 2.", logger="root") - super()._load_from_state_dict(state_dict, prefix, local_metadata, - strict, missing_keys, unexpected_keys, - error_msgs) + super()._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/multi_scale_deform_attn.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/multi_scale_deform_attn.py index c9129e22..7b2448b6 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/multi_scale_deform_attn.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/multi_scale_deform_attn.py @@ -7,21 +7,20 @@ import torch.nn.functional as F from torch.autograd.function import Function, once_differentiable -from comfy_controlnet_preprocessors.uniformer.mmcv import deprecated_api_warning -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import constant_init, xavier_init -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn.bricks.registry import ATTENTION -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import BaseModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv import deprecated_api_warning +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import constant_init, xavier_init +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn.bricks.registry import ATTENTION +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import BaseModule from ..utils import ext_loader -ext_module = ext_loader.load_ext( - '_ext', ['ms_deform_attn_backward', 'ms_deform_attn_forward']) +ext_module = ext_loader.load_ext("_ext", ["ms_deform_attn_backward", "ms_deform_attn_forward"]) class MultiScaleDeformableAttnFunction(Function): - @staticmethod - def forward(ctx, value, value_spatial_shapes, value_level_start_index, - sampling_locations, attention_weights, im2col_step): + def forward( + ctx, value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, im2col_step + ): """GPU version of multi-scale deformable attention. Args: @@ -50,10 +49,11 @@ def forward(ctx, value, value_spatial_shapes, value_level_start_index, value_level_start_index, sampling_locations, attention_weights, - im2col_step=ctx.im2col_step) - ctx.save_for_backward(value, value_spatial_shapes, - value_level_start_index, sampling_locations, - attention_weights) + im2col_step=ctx.im2col_step, + ) + ctx.save_for_backward( + value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights + ) return output @staticmethod @@ -69,8 +69,7 @@ def backward(ctx, grad_output): Tuple[Tensor]: Gradient of input tensors in forward. """ - value, value_spatial_shapes, value_level_start_index,\ - sampling_locations, attention_weights = ctx.saved_tensors + value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights = ctx.saved_tensors grad_value = torch.zeros_like(value) grad_sampling_loc = torch.zeros_like(sampling_locations) grad_attn_weight = torch.zeros_like(attention_weights) @@ -85,14 +84,13 @@ def backward(ctx, grad_output): grad_value, grad_sampling_loc, grad_attn_weight, - im2col_step=ctx.im2col_step) + im2col_step=ctx.im2col_step, + ) - return grad_value, None, None, \ - grad_sampling_loc, grad_attn_weight, None + return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None -def multi_scale_deformable_attn_pytorch(value, value_spatial_shapes, - sampling_locations, attention_weights): +def multi_scale_deformable_attn_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights): """CPU version of multi-scale deformable attention. Args: @@ -114,10 +112,8 @@ def multi_scale_deformable_attn_pytorch(value, value_spatial_shapes, """ bs, _, num_heads, embed_dims = value.shape - _, num_queries, num_heads, num_levels, num_points, _ =\ - sampling_locations.shape - value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], - dim=1) + _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape + value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1) sampling_grids = 2 * sampling_locations - 1 sampling_value_list = [] for level, (H_, W_) in enumerate(value_spatial_shapes): @@ -125,29 +121,27 @@ def multi_scale_deformable_attn_pytorch(value, value_spatial_shapes, # bs, H_*W_, num_heads*embed_dims -> # bs, num_heads*embed_dims, H_*W_ -> # bs*num_heads, embed_dims, H_, W_ - value_l_ = value_list[level].flatten(2).transpose(1, 2).reshape( - bs * num_heads, embed_dims, H_, W_) + value_l_ = value_list[level].flatten(2).transpose(1, 2).reshape(bs * num_heads, embed_dims, H_, W_) # bs, num_queries, num_heads, num_points, 2 -> # bs, num_heads, num_queries, num_points, 2 -> # bs*num_heads, num_queries, num_points, 2 - sampling_grid_l_ = sampling_grids[:, :, :, - level].transpose(1, 2).flatten(0, 1) + sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1) # bs*num_heads, embed_dims, num_queries, num_points sampling_value_l_ = F.grid_sample( - value_l_, - sampling_grid_l_, - mode='bilinear', - padding_mode='zeros', - align_corners=False) + value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False + ) sampling_value_list.append(sampling_value_l_) # (bs, num_queries, num_heads, num_levels, num_points) -> # (bs, num_heads, num_queries, num_levels, num_points) -> # (bs, num_heads, 1, num_queries, num_levels*num_points) attention_weights = attention_weights.transpose(1, 2).reshape( - bs * num_heads, 1, num_queries, num_levels * num_points) - output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * - attention_weights).sum(-1).view(bs, num_heads * embed_dims, - num_queries) + bs * num_heads, 1, num_queries, num_levels * num_points + ) + output = ( + (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights) + .sum(-1) + .view(bs, num_heads * embed_dims, num_queries) + ) return output.transpose(1, 2).contiguous() @@ -179,20 +173,21 @@ class MultiScaleDeformableAttention(BaseModule): Default: None. """ - def __init__(self, - embed_dims=256, - num_heads=8, - num_levels=4, - num_points=4, - im2col_step=64, - dropout=0.1, - batch_first=False, - norm_cfg=None, - init_cfg=None): + def __init__( + self, + embed_dims=256, + num_heads=8, + num_levels=4, + num_points=4, + im2col_step=64, + dropout=0.1, + batch_first=False, + norm_cfg=None, + init_cfg=None, + ): super().__init__(init_cfg) if embed_dims % num_heads != 0: - raise ValueError(f'embed_dims must be divisible by num_heads, ' - f'but got {embed_dims} and {num_heads}') + raise ValueError(f"embed_dims must be divisible by num_heads, " f"but got {embed_dims} and {num_heads}") dim_per_head = embed_dims // num_heads self.norm_cfg = norm_cfg self.dropout = nn.Dropout(dropout) @@ -202,64 +197,61 @@ def __init__(self, # which is more efficient in the CUDA implementation def _is_power_of_2(n): if (not isinstance(n, int)) or (n < 0): - raise ValueError( - 'invalid input for _is_power_of_2: {} (type: {})'.format( - n, type(n))) + raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n))) return (n & (n - 1) == 0) and n != 0 if not _is_power_of_2(dim_per_head): warnings.warn( "You'd better set embed_dims in " - 'MultiScaleDeformAttention to make ' - 'the dimension of each attention head a power of 2 ' - 'which is more efficient in our CUDA implementation.') + "MultiScaleDeformAttention to make " + "the dimension of each attention head a power of 2 " + "which is more efficient in our CUDA implementation." + ) self.im2col_step = im2col_step self.embed_dims = embed_dims self.num_levels = num_levels self.num_heads = num_heads self.num_points = num_points - self.sampling_offsets = nn.Linear( - embed_dims, num_heads * num_levels * num_points * 2) - self.attention_weights = nn.Linear(embed_dims, - num_heads * num_levels * num_points) + self.sampling_offsets = nn.Linear(embed_dims, num_heads * num_levels * num_points * 2) + self.attention_weights = nn.Linear(embed_dims, num_heads * num_levels * num_points) self.value_proj = nn.Linear(embed_dims, embed_dims) self.output_proj = nn.Linear(embed_dims, embed_dims) self.init_weights() def init_weights(self): """Default initialization for Parameters of Module.""" - constant_init(self.sampling_offsets, 0.) - thetas = torch.arange( - self.num_heads, - dtype=torch.float32) * (2.0 * math.pi / self.num_heads) + constant_init(self.sampling_offsets, 0.0) + thetas = torch.arange(self.num_heads, dtype=torch.float32) * (2.0 * math.pi / self.num_heads) grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) - grid_init = (grid_init / - grid_init.abs().max(-1, keepdim=True)[0]).view( - self.num_heads, 1, 1, - 2).repeat(1, self.num_levels, self.num_points, 1) + grid_init = ( + (grid_init / grid_init.abs().max(-1, keepdim=True)[0]) + .view(self.num_heads, 1, 1, 2) + .repeat(1, self.num_levels, self.num_points, 1) + ) for i in range(self.num_points): grid_init[:, :, i, :] *= i + 1 self.sampling_offsets.bias.data = grid_init.view(-1) - constant_init(self.attention_weights, val=0., bias=0.) - xavier_init(self.value_proj, distribution='uniform', bias=0.) - xavier_init(self.output_proj, distribution='uniform', bias=0.) + constant_init(self.attention_weights, val=0.0, bias=0.0) + xavier_init(self.value_proj, distribution="uniform", bias=0.0) + xavier_init(self.output_proj, distribution="uniform", bias=0.0) self._is_init = True - @deprecated_api_warning({'residual': 'identity'}, - cls_name='MultiScaleDeformableAttention') - def forward(self, - query, - key=None, - value=None, - identity=None, - query_pos=None, - key_padding_mask=None, - reference_points=None, - spatial_shapes=None, - level_start_index=None, - **kwargs): + @deprecated_api_warning({"residual": "identity"}, cls_name="MultiScaleDeformableAttention") + def forward( + self, + query, + key=None, + value=None, + identity=None, + query_pos=None, + key_padding_mask=None, + reference_points=None, + spatial_shapes=None, + level_start_index=None, + **kwargs, + ): """Forward Function of MultiScaleDeformAttention. Args: @@ -317,37 +309,35 @@ def forward(self, value = value.masked_fill(key_padding_mask[..., None], 0.0) value = value.view(bs, num_value, self.num_heads, -1) sampling_offsets = self.sampling_offsets(query).view( - bs, num_query, self.num_heads, self.num_levels, self.num_points, 2) + bs, num_query, self.num_heads, self.num_levels, self.num_points, 2 + ) attention_weights = self.attention_weights(query).view( - bs, num_query, self.num_heads, self.num_levels * self.num_points) + bs, num_query, self.num_heads, self.num_levels * self.num_points + ) attention_weights = attention_weights.softmax(-1) - attention_weights = attention_weights.view(bs, num_query, - self.num_heads, - self.num_levels, - self.num_points) + attention_weights = attention_weights.view(bs, num_query, self.num_heads, self.num_levels, self.num_points) if reference_points.shape[-1] == 2: - offset_normalizer = torch.stack( - [spatial_shapes[..., 1], spatial_shapes[..., 0]], -1) - sampling_locations = reference_points[:, :, None, :, None, :] \ - + sampling_offsets \ - / offset_normalizer[None, None, None, :, None, :] + offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1) + sampling_locations = ( + reference_points[:, :, None, :, None, :] + + sampling_offsets / offset_normalizer[None, None, None, :, None, :] + ) elif reference_points.shape[-1] == 4: - sampling_locations = reference_points[:, :, None, :, None, :2] \ - + sampling_offsets / self.num_points \ - * reference_points[:, :, None, :, None, 2:] \ - * 0.5 + sampling_locations = ( + reference_points[:, :, None, :, None, :2] + + sampling_offsets / self.num_points * reference_points[:, :, None, :, None, 2:] * 0.5 + ) else: raise ValueError( - f'Last dim of reference_points must be' - f' 2 or 4, but get {reference_points.shape[-1]} instead.') + f"Last dim of reference_points must be" f" 2 or 4, but get {reference_points.shape[-1]} instead." + ) if torch.cuda.is_available() and value.is_cuda: output = MultiScaleDeformableAttnFunction.apply( - value, spatial_shapes, level_start_index, sampling_locations, - attention_weights, self.im2col_step) + value, spatial_shapes, level_start_index, sampling_locations, attention_weights, self.im2col_step + ) else: - output = multi_scale_deformable_attn_pytorch( - value, spatial_shapes, sampling_locations, attention_weights) + output = multi_scale_deformable_attn_pytorch(value, spatial_shapes, sampling_locations, attention_weights) output = self.output_proj(output) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/nms.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/nms.py index b9791ef5..a03070b7 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/nms.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/nms.py @@ -3,28 +3,23 @@ import numpy as np import torch -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import deprecated_api_warning +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import deprecated_api_warning from ..utils import ext_loader -ext_module = ext_loader.load_ext( - '_ext', ['nms', 'softnms', 'nms_match', 'nms_rotated']) +ext_module = ext_loader.load_ext("_ext", ["nms", "softnms", "nms_match", "nms_rotated"]) # This function is modified from: https://github.com/pytorch/vision/ class NMSop(torch.autograd.Function): - @staticmethod - def forward(ctx, bboxes, scores, iou_threshold, offset, score_threshold, - max_num): + def forward(ctx, bboxes, scores, iou_threshold, offset, score_threshold, max_num): is_filtering_by_score = score_threshold > 0 if is_filtering_by_score: valid_mask = scores > score_threshold bboxes, scores = bboxes[valid_mask], scores[valid_mask] - valid_inds = torch.nonzero( - valid_mask, as_tuple=False).squeeze(dim=1) + valid_inds = torch.nonzero(valid_mask, as_tuple=False).squeeze(dim=1) - inds = ext_module.nms( - bboxes, scores, iou_threshold=float(iou_threshold), offset=offset) + inds = ext_module.nms(bboxes, scores, iou_threshold=float(iou_threshold), offset=offset) if max_num > 0: inds = inds[:max_num] @@ -33,19 +28,16 @@ def forward(ctx, bboxes, scores, iou_threshold, offset, score_threshold, return inds @staticmethod - def symbolic(g, bboxes, scores, iou_threshold, offset, score_threshold, - max_num): + def symbolic(g, bboxes, scores, iou_threshold, offset, score_threshold, max_num): from ..onnx import is_custom_op_loaded + has_custom_op = is_custom_op_loaded() # TensorRT nms plugin is aligned with original nms in ONNXRuntime - is_trt_backend = os.environ.get('ONNX_BACKEND') == 'MMCVTensorRT' + is_trt_backend = os.environ.get("ONNX_BACKEND") == "MMCVTensorRT" if has_custom_op and (not is_trt_backend): return g.op( - 'mmcv::NonMaxSuppression', - bboxes, - scores, - iou_threshold_f=float(iou_threshold), - offset_i=int(offset)) + "mmcv::NonMaxSuppression", bboxes, scores, iou_threshold_f=float(iou_threshold), offset_i=int(offset) + ) else: from torch.onnx.symbolic_opset9 import select, squeeze, unsqueeze from ..onnx.onnx_utils.symbolic_helper import _size_helper @@ -54,37 +46,21 @@ def symbolic(g, bboxes, scores, iou_threshold, offset, score_threshold, scores = unsqueeze(g, unsqueeze(g, scores, 0), 0) if max_num > 0: - max_num = g.op( - 'Constant', - value_t=torch.tensor(max_num, dtype=torch.long)) + max_num = g.op("Constant", value_t=torch.tensor(max_num, dtype=torch.long)) else: - dim = g.op('Constant', value_t=torch.tensor(0)) + dim = g.op("Constant", value_t=torch.tensor(0)) max_num = _size_helper(g, bboxes, dim) max_output_per_class = max_num - iou_threshold = g.op( - 'Constant', - value_t=torch.tensor([iou_threshold], dtype=torch.float)) - score_threshold = g.op( - 'Constant', - value_t=torch.tensor([score_threshold], dtype=torch.float)) - nms_out = g.op('NonMaxSuppression', boxes, scores, - max_output_per_class, iou_threshold, - score_threshold) - return squeeze( - g, - select( - g, nms_out, 1, - g.op( - 'Constant', - value_t=torch.tensor([2], dtype=torch.long))), 1) + iou_threshold = g.op("Constant", value_t=torch.tensor([iou_threshold], dtype=torch.float)) + score_threshold = g.op("Constant", value_t=torch.tensor([score_threshold], dtype=torch.float)) + nms_out = g.op("NonMaxSuppression", boxes, scores, max_output_per_class, iou_threshold, score_threshold) + return squeeze(g, select(g, nms_out, 1, g.op("Constant", value_t=torch.tensor([2], dtype=torch.long))), 1) class SoftNMSop(torch.autograd.Function): - @staticmethod - def forward(ctx, boxes, scores, iou_threshold, sigma, min_score, method, - offset): - dets = boxes.new_empty((boxes.size(0), 5), device='cpu') + def forward(ctx, boxes, scores, iou_threshold, sigma, min_score, method, offset): + dets = boxes.new_empty((boxes.size(0), 5), device="cpu") inds = ext_module.softnms( boxes.cpu(), scores.cpu(), @@ -93,16 +69,17 @@ def forward(ctx, boxes, scores, iou_threshold, sigma, min_score, method, sigma=float(sigma), min_score=float(min_score), method=int(method), - offset=int(offset)) + offset=int(offset), + ) return dets, inds @staticmethod - def symbolic(g, boxes, scores, iou_threshold, sigma, min_score, method, - offset): + def symbolic(g, boxes, scores, iou_threshold, sigma, min_score, method, offset): from packaging import version - assert version.parse(torch.__version__) >= version.parse('1.7.0') + + assert version.parse(torch.__version__) >= version.parse("1.7.0") nms_out = g.op( - 'mmcv::SoftNonMaxSuppression', + "mmcv::SoftNonMaxSuppression", boxes, scores, iou_threshold_f=float(iou_threshold), @@ -110,11 +87,12 @@ def symbolic(g, boxes, scores, iou_threshold, sigma, min_score, method, min_score_f=float(min_score), method_i=int(method), offset_i=int(offset), - outputs=2) + outputs=2, + ) return nms_out -@deprecated_api_warning({'iou_thr': 'iou_threshold'}) +@deprecated_api_warning({"iou_thr": "iou_threshold"}) def nms(boxes, scores, iou_threshold, offset=0, score_threshold=0, max_num=-1): """Dispatch to either CPU or GPU NMS implementations. @@ -160,16 +138,12 @@ def nms(boxes, scores, iou_threshold, offset=0, score_threshold=0, max_num=-1): assert boxes.size(0) == scores.size(0) assert offset in (0, 1) - if torch.__version__ == 'parrots': + if torch.__version__ == "parrots": indata_list = [boxes, scores] - indata_dict = { - 'iou_threshold': float(iou_threshold), - 'offset': int(offset) - } + indata_dict = {"iou_threshold": float(iou_threshold), "offset": int(offset)} inds = ext_module.nms(*indata_list, **indata_dict) else: - inds = NMSop.apply(boxes, scores, iou_threshold, offset, - score_threshold, max_num) + inds = NMSop.apply(boxes, scores, iou_threshold, offset, score_threshold, max_num) dets = torch.cat((boxes[inds], scores[inds].reshape(-1, 1)), dim=1) if is_numpy: dets = dets.cpu().numpy() @@ -177,14 +151,8 @@ def nms(boxes, scores, iou_threshold, offset=0, score_threshold=0, max_num=-1): return dets, inds -@deprecated_api_warning({'iou_thr': 'iou_threshold'}) -def soft_nms(boxes, - scores, - iou_threshold=0.3, - sigma=0.5, - min_score=1e-3, - method='linear', - offset=0): +@deprecated_api_warning({"iou_thr": "iou_threshold"}) +def soft_nms(boxes, scores, iou_threshold=0.3, sigma=0.5, min_score=1e-3, method="linear", offset=0): """Dispatch to only CPU Soft NMS implementations. The input can be either a torch tensor or numpy array. @@ -227,27 +195,32 @@ def soft_nms(boxes, assert boxes.size(1) == 4 assert boxes.size(0) == scores.size(0) assert offset in (0, 1) - method_dict = {'naive': 0, 'linear': 1, 'gaussian': 2} + method_dict = {"naive": 0, "linear": 1, "gaussian": 2} assert method in method_dict.keys() - if torch.__version__ == 'parrots': - dets = boxes.new_empty((boxes.size(0), 5), device='cpu') + if torch.__version__ == "parrots": + dets = boxes.new_empty((boxes.size(0), 5), device="cpu") indata_list = [boxes.cpu(), scores.cpu(), dets.cpu()] indata_dict = { - 'iou_threshold': float(iou_threshold), - 'sigma': float(sigma), - 'min_score': min_score, - 'method': method_dict[method], - 'offset': int(offset) + "iou_threshold": float(iou_threshold), + "sigma": float(sigma), + "min_score": min_score, + "method": method_dict[method], + "offset": int(offset), } inds = ext_module.softnms(*indata_list, **indata_dict) else: - dets, inds = SoftNMSop.apply(boxes.cpu(), scores.cpu(), - float(iou_threshold), float(sigma), - float(min_score), method_dict[method], - int(offset)) + dets, inds = SoftNMSop.apply( + boxes.cpu(), + scores.cpu(), + float(iou_threshold), + float(sigma), + float(min_score), + method_dict[method], + int(offset), + ) - dets = dets[:inds.size(0)] + dets = dets[: inds.size(0)] if is_numpy: dets = dets.cpu().numpy() @@ -290,7 +263,7 @@ def batched_nms(boxes, scores, idxs, nms_cfg, class_agnostic=False): tuple: kept dets and indice. """ nms_cfg_ = nms_cfg.copy() - class_agnostic = nms_cfg_.pop('class_agnostic', class_agnostic) + class_agnostic = nms_cfg_.pop("class_agnostic", class_agnostic) if class_agnostic: boxes_for_nms = boxes else: @@ -298,10 +271,10 @@ def batched_nms(boxes, scores, idxs, nms_cfg, class_agnostic=False): offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes)) boxes_for_nms = boxes + offsets[:, None] - nms_type = nms_cfg_.pop('type', 'nms') + nms_type = nms_cfg_.pop("type", "nms") nms_op = eval(nms_type) - split_thr = nms_cfg_.pop('split_thr', 10000) + split_thr = nms_cfg_.pop("split_thr", 10000) # Won't split to multiple nms nodes when exporting to onnx if boxes_for_nms.shape[0] < split_thr or torch.onnx.is_in_onnx_export(): dets, keep = nms_op(boxes_for_nms, scores, **nms_cfg_) @@ -313,7 +286,7 @@ def batched_nms(boxes, scores, idxs, nms_cfg, class_agnostic=False): # Some type of nms would reweight the score, such as SoftNMS scores = dets[:, 4] else: - max_num = nms_cfg_.pop('max_num', -1) + max_num = nms_cfg_.pop("max_num", -1) total_mask = scores.new_zeros(scores.size(), dtype=torch.bool) # Some type of nms would reweight the score, such as SoftNMS scores_after_nms = scores.new_zeros(scores.size()) @@ -355,16 +328,15 @@ def nms_match(dets, iou_threshold): if dets.shape[0] == 0: matched = [] else: - assert dets.shape[-1] == 5, 'inputs dets.shape should be (N, 5), ' \ - f'but get {dets.shape}' + assert dets.shape[-1] == 5, "inputs dets.shape should be (N, 5), " f"but get {dets.shape}" if isinstance(dets, torch.Tensor): dets_t = dets.detach().cpu() else: dets_t = torch.from_numpy(dets) indata_list = [dets_t] - indata_dict = {'iou_threshold': float(iou_threshold)} + indata_dict = {"iou_threshold": float(iou_threshold)} matched = ext_module.nms_match(*indata_list, **indata_dict) - if torch.__version__ == 'parrots': + if torch.__version__ == "parrots": matched = matched.tolist() if isinstance(dets, torch.Tensor): @@ -401,17 +373,11 @@ def nms_rotated(dets, scores, iou_threshold, labels=None): _, order = scores.sort(0, descending=True) dets_sorted = dets_wl.index_select(0, order) - if torch.__version__ == 'parrots': + if torch.__version__ == "parrots": keep_inds = ext_module.nms_rotated( - dets_wl, - scores, - order, - dets_sorted, - iou_threshold=iou_threshold, - multi_label=multi_label) + dets_wl, scores, order, dets_sorted, iou_threshold=iou_threshold, multi_label=multi_label + ) else: - keep_inds = ext_module.nms_rotated(dets_wl, scores, order, dets_sorted, - iou_threshold, multi_label) - dets = torch.cat((dets[keep_inds], scores[keep_inds].reshape(-1, 1)), - dim=1) + keep_inds = ext_module.nms_rotated(dets_wl, scores, order, dets_sorted, iou_threshold, multi_label) + dets = torch.cat((dets[keep_inds], scores[keep_inds].reshape(-1, 1)), dim=1) return dets, keep_inds diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/point_sample.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/point_sample.py index 743ca4a9..3e08f8c6 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/point_sample.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/point_sample.py @@ -53,7 +53,7 @@ def bilinear_grid_sample(im, grid, align_corners=False): wd = ((x - x0) * (y - y0)).unsqueeze(1) # Apply default for grid_sample function zero padding - im_padded = F.pad(im, pad=[1, 1, 1, 1], mode='constant', value=0) + im_padded = F.pad(im, pad=[1, 1, 1, 1], mode="constant", value=0) padded_h = h + 2 padded_w = w + 2 # save points positions after padding @@ -85,10 +85,10 @@ def bilinear_grid_sample(im, grid, align_corners=False): def is_in_onnx_export_without_custom_ops(): - from comfy_controlnet_preprocessors.uniformer.mmcv.ops import get_onnxruntime_op_path + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.ops import get_onnxruntime_op_path + ort_custom_op_path = get_onnxruntime_op_path() - return torch.onnx.is_in_onnx_export( - ) and not osp.exists(ort_custom_op_path) + return torch.onnx.is_in_onnx_export() and not osp.exists(ort_custom_op_path) def normalize(grid): @@ -127,9 +127,8 @@ def generate_grid(num_grid, size, device): contains coordinates for the regular grids. """ - affine_trans = torch.tensor([[[1., 0., 0.], [0., 1., 0.]]], device=device) - grid = F.affine_grid( - affine_trans, torch.Size((1, 1, *size)), align_corners=False) + affine_trans = torch.tensor([[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]], device=device) + grid = F.affine_grid(affine_trans, torch.Size((1, 1, *size)), align_corners=False) grid = normalize(grid) return grid.view(1, -1, 2).expand(num_grid, -1, -1) @@ -175,15 +174,13 @@ def get_shape_from_feature_map(x): torch.Tensor: Spatial resolution (width, height), shape (1, 1, 2) """ if torch.onnx.is_in_onnx_export(): - img_shape = shape_as_tensor(x)[2:].flip(0).view(1, 1, 2).to( - x.device).float() + img_shape = shape_as_tensor(x)[2:].flip(0).view(1, 1, 2).to(x.device).float() else: - img_shape = torch.tensor(x.shape[2:]).flip(0).view(1, 1, 2).to( - x.device).float() + img_shape = torch.tensor(x.shape[2:]).flip(0).view(1, 1, 2).to(x.device).float() return img_shape -def abs_img_point_to_rel_img_point(abs_img_points, img, spatial_scale=1.): +def abs_img_point_to_rel_img_point(abs_img_points, img, spatial_scale=1.0): """Convert image based absolute point coordinates to image based relative coordinates for sampling. @@ -198,14 +195,11 @@ def abs_img_point_to_rel_img_point(abs_img_points, img, spatial_scale=1.): shape (N, P, 2) """ - assert (isinstance(img, tuple) and len(img) == 2) or \ - (isinstance(img, torch.Tensor) and len(img.shape) == 4) + assert (isinstance(img, tuple) and len(img) == 2) or (isinstance(img, torch.Tensor) and len(img.shape) == 4) if isinstance(img, tuple): h, w = img - scale = torch.tensor([w, h], - dtype=torch.float, - device=abs_img_points.device) + scale = torch.tensor([w, h], dtype=torch.float, device=abs_img_points.device) scale = scale.view(1, 1, 2) else: scale = get_shape_from_feature_map(img) @@ -213,10 +207,7 @@ def abs_img_point_to_rel_img_point(abs_img_points, img, spatial_scale=1.): return abs_img_points / scale * spatial_scale -def rel_roi_point_to_rel_img_point(rois, - rel_roi_points, - img, - spatial_scale=1.): +def rel_roi_point_to_rel_img_point(rois, rel_roi_points, img, spatial_scale=1.0): """Convert roi based relative point coordinates to image based absolute point coordinates. @@ -233,8 +224,7 @@ def rel_roi_point_to_rel_img_point(rois, """ abs_img_point = rel_roi_point_to_abs_img_point(rois, rel_roi_points) - rel_img_point = abs_img_point_to_rel_img_point(abs_img_point, img, - spatial_scale) + rel_img_point = abs_img_point_to_rel_img_point(abs_img_point, img, spatial_scale) return rel_img_point @@ -263,18 +253,15 @@ def point_sample(input, points, align_corners=False, **kwargs): # If custom ops for onnx runtime not compiled use python # implementation of grid_sample function to make onnx graph # with supported nodes - output = bilinear_grid_sample( - input, denormalize(points), align_corners=align_corners) + output = bilinear_grid_sample(input, denormalize(points), align_corners=align_corners) else: - output = F.grid_sample( - input, denormalize(points), align_corners=align_corners, **kwargs) + output = F.grid_sample(input, denormalize(points), align_corners=align_corners, **kwargs) if add_dim: output = output.squeeze(3) return output class SimpleRoIAlign(nn.Module): - def __init__(self, output_size, spatial_scale, aligned=True): """Simple RoI align in PointRend, faster than standard RoIAlign. @@ -296,29 +283,24 @@ def __init__(self, output_size, spatial_scale, aligned=True): def forward(self, features, rois): num_imgs = features.size(0) num_rois = rois.size(0) - rel_roi_points = generate_grid( - num_rois, self.output_size, device=rois.device) + rel_roi_points = generate_grid(num_rois, self.output_size, device=rois.device) if torch.onnx.is_in_onnx_export(): - rel_img_points = rel_roi_point_to_rel_img_point( - rois, rel_roi_points, features, self.spatial_scale) - rel_img_points = rel_img_points.reshape(num_imgs, -1, - *rel_img_points.shape[1:]) - point_feats = point_sample( - features, rel_img_points, align_corners=not self.aligned) + rel_img_points = rel_roi_point_to_rel_img_point(rois, rel_roi_points, features, self.spatial_scale) + rel_img_points = rel_img_points.reshape(num_imgs, -1, *rel_img_points.shape[1:]) + point_feats = point_sample(features, rel_img_points, align_corners=not self.aligned) point_feats = point_feats.transpose(1, 2) else: point_feats = [] for batch_ind in range(num_imgs): # unravel batch dim feat = features[batch_ind].unsqueeze(0) - inds = (rois[:, 0].long() == batch_ind) + inds = rois[:, 0].long() == batch_ind if inds.any(): rel_img_points = rel_roi_point_to_rel_img_point( - rois[inds], rel_roi_points[inds], feat, - self.spatial_scale).unsqueeze(0) - point_feat = point_sample( - feat, rel_img_points, align_corners=not self.aligned) + rois[inds], rel_roi_points[inds], feat, self.spatial_scale + ).unsqueeze(0) + point_feat = point_sample(feat, rel_img_points, align_corners=not self.aligned) point_feat = point_feat.squeeze(0).transpose(0, 1) point_feats.append(point_feat) @@ -331,6 +313,5 @@ def forward(self, features, rois): def __repr__(self): format_str = self.__class__.__name__ - format_str += '(output_size={}, spatial_scale={}'.format( - self.output_size, self.spatial_scale) + format_str += "(output_size={}, spatial_scale={}".format(self.output_size, self.spatial_scale) return format_str diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/points_sampler.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/points_sampler.py index 3edf1ef0..4286ed45 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/points_sampler.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/points_sampler.py @@ -3,9 +3,8 @@ import torch from torch import nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import force_fp32 -from .furthest_point_sample import (furthest_point_sample, - furthest_point_sample_with_dist) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import force_fp32 +from .furthest_point_sample import furthest_point_sample, furthest_point_sample_with_dist def calc_square_dist(point_feat_a, point_feat_b, norm=True): @@ -45,16 +44,17 @@ def get_sampler_cls(sampler_type): class: Points sampler type. """ sampler_mappings = { - 'D-FPS': DFPSSampler, - 'F-FPS': FFPSSampler, - 'FS': FSSampler, + "D-FPS": DFPSSampler, + "F-FPS": FFPSSampler, + "FS": FSSampler, } try: return sampler_mappings[sampler_type] except KeyError: raise KeyError( - f'Supported `sampler_type` are {sampler_mappings.keys()}, but got \ - {sampler_type}') + f"Supported `sampler_type` are {sampler_mappings.keys()}, but got \ + {sampler_type}" + ) class PointsSampler(nn.Module): @@ -71,16 +71,14 @@ class PointsSampler(nn.Module): Range of points to apply FPS. Default: [-1]. """ - def __init__(self, - num_point: List[int], - fps_mod_list: List[str] = ['D-FPS'], - fps_sample_range_list: List[int] = [-1]): + def __init__( + self, num_point: List[int], fps_mod_list: List[str] = ["D-FPS"], fps_sample_range_list: List[int] = [-1] + ): super().__init__() # FPS would be applied to different fps_mod in the list, # so the length of the num_point should be equal to # fps_mod_list and fps_sample_range_list. - assert len(num_point) == len(fps_mod_list) == len( - fps_sample_range_list) + assert len(num_point) == len(fps_mod_list) == len(fps_sample_range_list) self.num_point = num_point self.fps_sample_range_list = fps_sample_range_list self.samplers = nn.ModuleList() @@ -101,8 +99,7 @@ def forward(self, points_xyz, features): indices = [] last_fps_end_index = 0 - for fps_sample_range, sampler, npoint in zip( - self.fps_sample_range_list, self.samplers, self.num_point): + for fps_sample_range, sampler, npoint in zip(self.fps_sample_range_list, self.samplers, self.num_point): assert fps_sample_range < points_xyz.shape[1] if fps_sample_range == -1: @@ -112,16 +109,13 @@ def forward(self, points_xyz, features): else: sample_features = None else: - sample_points_xyz = \ - points_xyz[:, last_fps_end_index:fps_sample_range] + sample_points_xyz = points_xyz[:, last_fps_end_index:fps_sample_range] if features is not None: - sample_features = features[:, :, last_fps_end_index: - fps_sample_range] + sample_features = features[:, :, last_fps_end_index:fps_sample_range] else: sample_features = None - fps_idx = sampler(sample_points_xyz.contiguous(), sample_features, - npoint) + fps_idx = sampler(sample_points_xyz.contiguous(), sample_features, npoint) indices.append(fps_idx + last_fps_end_index) last_fps_end_index += fps_sample_range @@ -150,11 +144,9 @@ def __init__(self): def forward(self, points, features, npoint): """Sampling points with F-FPS.""" - assert features is not None, \ - 'feature input to FFPS_Sampler should not be None' + assert features is not None, "feature input to FFPS_Sampler should not be None" features_for_fps = torch.cat([points, features.transpose(1, 2)], dim=2) - features_dist = calc_square_dist( - features_for_fps, features_for_fps, norm=False) + features_dist = calc_square_dist(features_for_fps, features_for_fps, norm=False) fps_idx = furthest_point_sample_with_dist(features_dist, npoint) return fps_idx @@ -167,8 +159,7 @@ def __init__(self): def forward(self, points, features, npoint): """Sampling points with FS_Sampling.""" - assert features is not None, \ - 'feature input to FS_Sampler should not be None' + assert features is not None, "feature input to FS_Sampler should not be None" ffps_sampler = FFPSSampler() dfps_sampler = DFPSSampler() fps_idx_ffps = ffps_sampler(points, features, npoint) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/roiaware_pool3d.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/roiaware_pool3d.py index d67a336e..76013012 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/roiaware_pool3d.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/roiaware_pool3d.py @@ -3,11 +3,10 @@ from torch import nn as nn from torch.autograd import Function -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv from ..utils import ext_loader -ext_module = ext_loader.load_ext( - '_ext', ['roiaware_pool3d_forward', 'roiaware_pool3d_backward']) +ext_module = ext_loader.load_ext("_ext", ["roiaware_pool3d_forward", "roiaware_pool3d_backward"]) class RoIAwarePool3d(nn.Module): @@ -25,13 +24,13 @@ class RoIAwarePool3d(nn.Module): Default: 'max'. """ - def __init__(self, out_size, max_pts_per_voxel=128, mode='max'): + def __init__(self, out_size, max_pts_per_voxel=128, mode="max"): super().__init__() self.out_size = out_size self.max_pts_per_voxel = max_pts_per_voxel - assert mode in ['max', 'avg'] - pool_mapping = {'max': 0, 'avg': 1} + assert mode in ["max", "avg"] + pool_mapping = {"max": 0, "avg": 1} self.mode = pool_mapping[mode] def forward(self, rois, pts, pts_feature): @@ -46,16 +45,12 @@ def forward(self, rois, pts, pts_feature): pooled_features (torch.Tensor): [N, out_x, out_y, out_z, C] """ - return RoIAwarePool3dFunction.apply(rois, pts, pts_feature, - self.out_size, - self.max_pts_per_voxel, self.mode) + return RoIAwarePool3dFunction.apply(rois, pts, pts_feature, self.out_size, self.max_pts_per_voxel, self.mode) class RoIAwarePool3dFunction(Function): - @staticmethod - def forward(ctx, rois, pts, pts_feature, out_size, max_pts_per_voxel, - mode): + def forward(ctx, rois, pts, pts_feature, out_size, max_pts_per_voxel, mode): """ Args: rois (torch.Tensor): [N, 7], in LiDAR coordinate, @@ -85,20 +80,13 @@ def forward(ctx, rois, pts, pts_feature, out_size, max_pts_per_voxel, num_channels = pts_feature.shape[-1] num_pts = pts.shape[0] - pooled_features = pts_feature.new_zeros( - (num_rois, out_x, out_y, out_z, num_channels)) - argmax = pts_feature.new_zeros( - (num_rois, out_x, out_y, out_z, num_channels), dtype=torch.int) - pts_idx_of_voxels = pts_feature.new_zeros( - (num_rois, out_x, out_y, out_z, max_pts_per_voxel), - dtype=torch.int) + pooled_features = pts_feature.new_zeros((num_rois, out_x, out_y, out_z, num_channels)) + argmax = pts_feature.new_zeros((num_rois, out_x, out_y, out_z, num_channels), dtype=torch.int) + pts_idx_of_voxels = pts_feature.new_zeros((num_rois, out_x, out_y, out_z, max_pts_per_voxel), dtype=torch.int) - ext_module.roiaware_pool3d_forward(rois, pts, pts_feature, argmax, - pts_idx_of_voxels, pooled_features, - mode) + ext_module.roiaware_pool3d_forward(rois, pts, pts_feature, argmax, pts_idx_of_voxels, pooled_features, mode) - ctx.roiaware_pool3d_for_backward = (pts_idx_of_voxels, argmax, mode, - num_pts, num_channels) + ctx.roiaware_pool3d_for_backward = (pts_idx_of_voxels, argmax, mode, num_pts, num_channels) return pooled_features @staticmethod @@ -107,8 +95,6 @@ def backward(ctx, grad_out): pts_idx_of_voxels, argmax, mode, num_pts, num_channels = ret grad_in = grad_out.new_zeros((num_pts, num_channels)) - ext_module.roiaware_pool3d_backward(pts_idx_of_voxels, argmax, - grad_out.contiguous(), grad_in, - mode) + ext_module.roiaware_pool3d_backward(pts_idx_of_voxels, argmax, grad_out.contiguous(), grad_in, mode) return None, None, grad_in, None, None, None diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/saconv.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/saconv.py index 350d392a..5e7c02c7 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/saconv.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/saconv.py @@ -3,12 +3,12 @@ import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import CONV_LAYERS, ConvAWS2d, constant_init -from comfy_controlnet_preprocessors.uniformer.mmcv.ops.deform_conv import deform_conv2d -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, digit_version +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import CONV_LAYERS, ConvAWS2d, constant_init +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.ops.deform_conv import deform_conv2d +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, digit_version -@CONV_LAYERS.register_module(name='SAC') +@CONV_LAYERS.register_module(name="SAC") class SAConv2d(ConvAWS2d): """SAC (Switchable Atrous Convolution) @@ -34,16 +34,18 @@ class SAConv2d(ConvAWS2d): convolution. Default: ``False``. """ - def __init__(self, - in_channels, - out_channels, - kernel_size, - stride=1, - padding=0, - dilation=1, - groups=1, - bias=True, - use_deform=False): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + groups=1, + bias=True, + use_deform=False, + ): super().__init__( in_channels, out_channels, @@ -52,30 +54,16 @@ def __init__(self, padding=padding, dilation=dilation, groups=groups, - bias=bias) + bias=bias, + ) self.use_deform = use_deform - self.switch = nn.Conv2d( - self.in_channels, 1, kernel_size=1, stride=stride, bias=True) + self.switch = nn.Conv2d(self.in_channels, 1, kernel_size=1, stride=stride, bias=True) self.weight_diff = nn.Parameter(torch.Tensor(self.weight.size())) - self.pre_context = nn.Conv2d( - self.in_channels, self.in_channels, kernel_size=1, bias=True) - self.post_context = nn.Conv2d( - self.out_channels, self.out_channels, kernel_size=1, bias=True) + self.pre_context = nn.Conv2d(self.in_channels, self.in_channels, kernel_size=1, bias=True) + self.post_context = nn.Conv2d(self.out_channels, self.out_channels, kernel_size=1, bias=True) if self.use_deform: - self.offset_s = nn.Conv2d( - self.in_channels, - 18, - kernel_size=3, - padding=1, - stride=stride, - bias=True) - self.offset_l = nn.Conv2d( - self.in_channels, - 18, - kernel_size=3, - padding=1, - stride=stride, - bias=True) + self.offset_s = nn.Conv2d(self.in_channels, 18, kernel_size=3, padding=1, stride=stride, bias=True) + self.offset_l = nn.Conv2d(self.in_channels, 18, kernel_size=3, padding=1, stride=stride, bias=True) self.init_weights() def init_weights(self): @@ -94,23 +82,20 @@ def forward(self, x): avg_x = avg_x.expand_as(x) x = x + avg_x # switch - avg_x = F.pad(x, pad=(2, 2, 2, 2), mode='reflect') + avg_x = F.pad(x, pad=(2, 2, 2, 2), mode="reflect") avg_x = F.avg_pool2d(avg_x, kernel_size=5, stride=1, padding=0) switch = self.switch(avg_x) # sac weight = self._get_weight(self.weight) - zero_bias = torch.zeros( - self.out_channels, device=weight.device, dtype=weight.dtype) + zero_bias = torch.zeros(self.out_channels, device=weight.device, dtype=weight.dtype) if self.use_deform: offset = self.offset_s(avg_x) - out_s = deform_conv2d(x, offset, weight, self.stride, self.padding, - self.dilation, self.groups, 1) + out_s = deform_conv2d(x, offset, weight, self.stride, self.padding, self.dilation, self.groups, 1) else: - if (TORCH_VERSION == 'parrots' - or digit_version(TORCH_VERSION) < digit_version('1.5.0')): + if TORCH_VERSION == "parrots" or digit_version(TORCH_VERSION) < digit_version("1.5.0"): out_s = super().conv2d_forward(x, weight) - elif digit_version(TORCH_VERSION) >= digit_version('1.8.0'): + elif digit_version(TORCH_VERSION) >= digit_version("1.8.0"): # bias is a required argument of _conv_forward in torch 1.8.0 out_s = super()._conv_forward(x, weight, zero_bias) else: @@ -122,13 +107,11 @@ def forward(self, x): weight = weight + self.weight_diff if self.use_deform: offset = self.offset_l(avg_x) - out_l = deform_conv2d(x, offset, weight, self.stride, self.padding, - self.dilation, self.groups, 1) + out_l = deform_conv2d(x, offset, weight, self.stride, self.padding, self.dilation, self.groups, 1) else: - if (TORCH_VERSION == 'parrots' - or digit_version(TORCH_VERSION) < digit_version('1.5.0')): + if TORCH_VERSION == "parrots" or digit_version(TORCH_VERSION) < digit_version("1.5.0"): out_l = super().conv2d_forward(x, weight) - elif digit_version(TORCH_VERSION) >= digit_version('1.8.0'): + elif digit_version(TORCH_VERSION) >= digit_version("1.8.0"): # bias is a required argument of _conv_forward in torch 1.8.0 out_l = super()._conv_forward(x, weight, zero_bias) else: diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/sync_bn.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/sync_bn.py index e8872f49..1cc3c122 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/sync_bn.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/sync_bn.py @@ -7,22 +7,26 @@ from torch.nn.modules.module import Module from torch.nn.parameter import Parameter -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import NORM_LAYERS +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import NORM_LAYERS from ..utils import ext_loader -ext_module = ext_loader.load_ext('_ext', [ - 'sync_bn_forward_mean', 'sync_bn_forward_var', 'sync_bn_forward_output', - 'sync_bn_backward_param', 'sync_bn_backward_data' -]) +ext_module = ext_loader.load_ext( + "_ext", + [ + "sync_bn_forward_mean", + "sync_bn_forward_var", + "sync_bn_forward_output", + "sync_bn_backward_param", + "sync_bn_backward_data", + ], +) class SyncBatchNormFunction(Function): - @staticmethod - def symbolic(g, input, running_mean, running_var, weight, bias, momentum, - eps, group, group_size, stats_mode): + def symbolic(g, input, running_mean, running_var, weight, bias, momentum, eps, group, group_size, stats_mode): return g.op( - 'mmcv::MMCVSyncBatchNorm', + "mmcv::MMCVSyncBatchNorm", input, running_mean, running_var, @@ -32,11 +36,11 @@ def symbolic(g, input, running_mean, running_var, weight, bias, momentum, eps_f=eps, group_i=group, group_size_i=group_size, - stats_mode=stats_mode) + stats_mode=stats_mode, + ) @staticmethod - def forward(self, input, running_mean, running_var, weight, bias, momentum, - eps, group, group_size, stats_mode): + def forward(self, input, running_mean, running_var, weight, bias, momentum, eps, group, group_size, stats_mode): self.momentum = momentum self.eps = eps self.group = group @@ -44,9 +48,8 @@ def forward(self, input, running_mean, running_var, weight, bias, momentum, self.stats_mode = stats_mode assert isinstance( - input, (torch.HalfTensor, torch.FloatTensor, - torch.cuda.HalfTensor, torch.cuda.FloatTensor)), \ - f'only support Half or Float Tensor, but {input.type()}' + input, (torch.HalfTensor, torch.FloatTensor, torch.cuda.HalfTensor, torch.cuda.FloatTensor) + ), f"only support Half or Float Tensor, but {input.type()}" output = torch.zeros_like(input) input3d = input.flatten(start_dim=2) output3d = output.view_as(input3d) @@ -54,14 +57,10 @@ def forward(self, input, running_mean, running_var, weight, bias, momentum, # ensure mean/var/norm/std are initialized as zeros # ``torch.empty()`` does not guarantee that - mean = torch.zeros( - num_channels, dtype=torch.float, device=input3d.device) - var = torch.zeros( - num_channels, dtype=torch.float, device=input3d.device) - norm = torch.zeros_like( - input3d, dtype=torch.float, device=input3d.device) - std = torch.zeros( - num_channels, dtype=torch.float, device=input3d.device) + mean = torch.zeros(num_channels, dtype=torch.float, device=input3d.device) + var = torch.zeros(num_channels, dtype=torch.float, device=input3d.device) + norm = torch.zeros_like(input3d, dtype=torch.float, device=input3d.device) + std = torch.zeros(num_channels, dtype=torch.float, device=input3d.device) batch_size = input3d.size(0) if batch_size > 0: @@ -73,16 +72,16 @@ def forward(self, input, running_mean, running_var, weight, bias, momentum, # synchronize mean and the batch flag vec = torch.cat([mean, batch_flag]) - if self.stats_mode == 'N': + if self.stats_mode == "N": vec *= batch_size if self.group_size > 1: dist.all_reduce(vec, group=self.group) total_batch = vec[-1].detach() mean = vec[:num_channels] - if self.stats_mode == 'default': + if self.stats_mode == "default": mean = mean / self.group_size - elif self.stats_mode == 'N': + elif self.stats_mode == "N": mean = mean / total_batch.clamp(min=1) else: raise NotImplementedError @@ -91,14 +90,14 @@ def forward(self, input, running_mean, running_var, weight, bias, momentum, if batch_size > 0: ext_module.sync_bn_forward_var(input3d, mean, var) - if self.stats_mode == 'N': + if self.stats_mode == "N": var *= batch_size if self.group_size > 1: dist.all_reduce(var, group=self.group) - if self.stats_mode == 'default': + if self.stats_mode == "default": var /= self.group_size - elif self.stats_mode == 'N': + elif self.stats_mode == "N": var /= total_batch.clamp(min=1) else: raise NotImplementedError @@ -120,7 +119,8 @@ def forward(self, input, running_mean, running_var, weight, bias, momentum, output3d, eps=self.eps, momentum=momentum, - group_size=self.group_size) + group_size=self.group_size, + ) self.save_for_backward(norm, std, weight) return output @@ -136,8 +136,7 @@ def backward(self, grad_output): batch_size = grad_input3d.size(0) if batch_size > 0: - ext_module.sync_bn_backward_param(grad_output3d, norm, grad_weight, - grad_bias) + ext_module.sync_bn_backward_param(grad_output3d, norm, grad_weight, grad_bias) # all reduce if self.group_size > 1: @@ -147,15 +146,12 @@ def backward(self, grad_output): grad_bias /= self.group_size if batch_size > 0: - ext_module.sync_bn_backward_data(grad_output3d, weight, - grad_weight, grad_bias, norm, std, - grad_input3d) + ext_module.sync_bn_backward_data(grad_output3d, weight, grad_weight, grad_bias, norm, std, grad_input3d) - return grad_input, None, None, grad_weight, grad_bias, \ - None, None, None, None, None + return grad_input, None, None, grad_weight, grad_bias, None, None, None, None, None -@NORM_LAYERS.register_module(name='MMSyncBN') +@NORM_LAYERS.register_module(name="MMSyncBN") class SyncBatchNorm(Module): """Synchronized Batch Normalization. @@ -190,14 +186,16 @@ class SyncBatchNorm(Module): number of batch. """ - def __init__(self, - num_features, - eps=1e-5, - momentum=0.1, - affine=True, - track_running_stats=True, - group=None, - stats_mode='default'): + def __init__( + self, + num_features, + eps=1e-5, + momentum=0.1, + affine=True, + track_running_stats=True, + group=None, + stats_mode="default", + ): super(SyncBatchNorm, self).__init__() self.num_features = num_features self.eps = eps @@ -207,24 +205,22 @@ def __init__(self, group = dist.group.WORLD if group is None else group self.group = group self.group_size = dist.get_world_size(group) - assert stats_mode in ['default', 'N'], \ - f'"stats_mode" only accepts "default" and "N", got "{stats_mode}"' + assert stats_mode in ["default", "N"], f'"stats_mode" only accepts "default" and "N", got "{stats_mode}"' self.stats_mode = stats_mode if self.affine: self.weight = Parameter(torch.Tensor(num_features)) self.bias = Parameter(torch.Tensor(num_features)) else: - self.register_parameter('weight', None) - self.register_parameter('bias', None) + self.register_parameter("weight", None) + self.register_parameter("bias", None) if self.track_running_stats: - self.register_buffer('running_mean', torch.zeros(num_features)) - self.register_buffer('running_var', torch.ones(num_features)) - self.register_buffer('num_batches_tracked', - torch.tensor(0, dtype=torch.long)) + self.register_buffer("running_mean", torch.zeros(num_features)) + self.register_buffer("running_var", torch.ones(num_features)) + self.register_buffer("num_batches_tracked", torch.tensor(0, dtype=torch.long)) else: - self.register_buffer('running_mean', None) - self.register_buffer('running_var', None) - self.register_buffer('num_batches_tracked', None) + self.register_buffer("running_mean", None) + self.register_buffer("running_var", None) + self.register_buffer("num_batches_tracked", None) self.reset_parameters() def reset_running_stats(self): @@ -241,8 +237,7 @@ def reset_parameters(self): def forward(self, input): if input.dim() < 2: - raise ValueError( - f'expected at least 2D input, got {input.dim()}D input') + raise ValueError(f"expected at least 2D input, got {input.dim()}D input") if self.momentum is None: exponential_average_factor = 0.0 else: @@ -252,28 +247,42 @@ def forward(self, input): if self.num_batches_tracked is not None: self.num_batches_tracked += 1 if self.momentum is None: # use cumulative moving average - exponential_average_factor = 1.0 / float( - self.num_batches_tracked) + exponential_average_factor = 1.0 / float(self.num_batches_tracked) else: # use exponential moving average exponential_average_factor = self.momentum if self.training or not self.track_running_stats: return SyncBatchNormFunction.apply( - input, self.running_mean, self.running_var, self.weight, - self.bias, exponential_average_factor, self.eps, self.group, - self.group_size, self.stats_mode) + input, + self.running_mean, + self.running_var, + self.weight, + self.bias, + exponential_average_factor, + self.eps, + self.group, + self.group_size, + self.stats_mode, + ) else: - return F.batch_norm(input, self.running_mean, self.running_var, - self.weight, self.bias, False, - exponential_average_factor, self.eps) + return F.batch_norm( + input, + self.running_mean, + self.running_var, + self.weight, + self.bias, + False, + exponential_average_factor, + self.eps, + ) def __repr__(self): s = self.__class__.__name__ - s += f'({self.num_features}, ' - s += f'eps={self.eps}, ' - s += f'momentum={self.momentum}, ' - s += f'affine={self.affine}, ' - s += f'track_running_stats={self.track_running_stats}, ' - s += f'group_size={self.group_size},' - s += f'stats_mode={self.stats_mode})' + s += f"({self.num_features}, " + s += f"eps={self.eps}, " + s += f"momentum={self.momentum}, " + s += f"affine={self.affine}, " + s += f"track_running_stats={self.track_running_stats}, " + s += f"group_size={self.group_size}," + s += f"stats_mode={self.stats_mode})" return s diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/upfirdn2d.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/upfirdn2d.py index 03b6a984..d1161baf 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/upfirdn2d.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/ops/upfirdn2d.py @@ -99,18 +99,15 @@ from torch.autograd import Function from torch.nn import functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import to_2tuple +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import to_2tuple from ..utils import ext_loader -upfirdn2d_ext = ext_loader.load_ext('_ext', ['upfirdn2d']) +upfirdn2d_ext = ext_loader.load_ext("_ext", ["upfirdn2d"]) class UpFirDn2dBackward(Function): - @staticmethod - def forward(ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, - in_size, out_size): - + def forward(ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size): up_x, up_y = up down_x, down_y = down g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad @@ -127,9 +124,9 @@ def forward(ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, pad_x0=g_pad_x0, pad_x1=g_pad_x1, pad_y0=g_pad_y0, - pad_y1=g_pad_y1) - grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], - in_size[3]) + pad_y1=g_pad_y1, + ) + grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3]) ctx.save_for_backward(kernel) @@ -150,10 +147,9 @@ def forward(ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, @staticmethod def backward(ctx, gradgrad_input): - kernel, = ctx.saved_tensors + (kernel,) = ctx.saved_tensors - gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], - ctx.in_size[3], 1) + gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1) gradgrad_out = upfirdn2d_ext.upfirdn2d( gradgrad_input, @@ -165,17 +161,16 @@ def backward(ctx, gradgrad_input): pad_x0=ctx.pad_x0, pad_x1=ctx.pad_x1, pad_y0=ctx.pad_y0, - pad_y1=ctx.pad_y1) + pad_y1=ctx.pad_y1, + ) # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], # ctx.out_size[1], ctx.in_size[3]) - gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.in_size[1], - ctx.out_size[0], ctx.out_size[1]) + gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]) return gradgrad_out, None, None, None, None, None, None, None, None class UpFirDn2d(Function): - @staticmethod def forward(ctx, input, kernel, up, down, pad): up_x, up_y = up @@ -215,7 +210,8 @@ def forward(ctx, input, kernel, up, down, pad): pad_x0=pad_x0, pad_x1=pad_x1, pad_y0=pad_y0, - pad_y1=pad_y1) + pad_y1=pad_y1, + ) # out = out.view(major, out_h, out_w, minor) out = out.view(-1, channel, out_h, out_w) @@ -262,7 +258,7 @@ def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): Returns: Tensor: Tensor after UpFIRDn. """ - if input.device.type == 'cpu': + if input.device.type == "cpu": if len(pad) == 2: pad = (pad[0], pad[1], pad[0], pad[1]) @@ -270,8 +266,7 @@ def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): down = to_2tuple(down) - out = upfirdn2d_native(input, kernel, up[0], up[1], down[0], down[1], - pad[0], pad[1], pad[2], pad[3]) + out = upfirdn2d_native(input, kernel, up[0], up[1], down[0], down[1], pad[0], pad[1], pad[2], pad[3]) else: _up = to_2tuple(up) @@ -287,8 +282,7 @@ def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): return out -def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, - pad_y0, pad_y1): +def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1): _, channel, in_h, in_w = input.shape input = input.reshape(-1, in_h, in_w, 1) @@ -299,20 +293,16 @@ def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1]) out = out.view(-1, in_h * up_y, in_w * up_x, minor) - out = F.pad( - out, - [0, 0, - max(pad_x0, 0), - max(pad_x1, 0), - max(pad_y0, 0), - max(pad_y1, 0)]) - out = out[:, - max(-pad_y0, 0):out.shape[1] - max(-pad_y1, 0), - max(-pad_x0, 0):out.shape[2] - max(-pad_x1, 0), :, ] + out = F.pad(out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]) + out = out[ + :, + max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0), + max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0), + :, + ] out = out.permute(0, 3, 1, 2) - out = out.reshape( - [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]) + out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]) w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w) out = F.conv2d(out, w) out = out.reshape( diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/distributed.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/distributed.py index 26068184..30fa7c7a 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/distributed.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/distributed.py @@ -1,10 +1,9 @@ # Copyright (c) OpenMMLab. All rights reserved. import torch -from torch.nn.parallel.distributed import (DistributedDataParallel, - _find_tensors) +from torch.nn.parallel.distributed import DistributedDataParallel, _find_tensors -from comfy_controlnet_preprocessors.uniformer.mmcv import print_log -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, digit_version +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv import print_log +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, digit_version from .scatter_gather import scatter_kwargs @@ -37,35 +36,32 @@ def train_step(self, *inputs, **kwargs): # In PyTorch >= 1.7, ``reducer._rebuild_buckets()`` is moved from the # end of backward to the beginning of forward. - if ('parrots' not in TORCH_VERSION - and digit_version(TORCH_VERSION) >= digit_version('1.7') - and self.reducer._rebuild_buckets()): - print_log( - 'Reducer buckets have been rebuilt in this iteration.', - logger='mmcv') - - if getattr(self, 'require_forward_param_sync', True): + if ( + "parrots" not in TORCH_VERSION + and digit_version(TORCH_VERSION) >= digit_version("1.7") + and self.reducer._rebuild_buckets() + ): + print_log("Reducer buckets have been rebuilt in this iteration.", logger="mmcv") + + if getattr(self, "require_forward_param_sync", True): self._sync_params() if self.device_ids: inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids) if len(self.device_ids) == 1: output = self.module.train_step(*inputs[0], **kwargs[0]) else: - outputs = self.parallel_apply( - self._module_copies[:len(inputs)], inputs, kwargs) + outputs = self.parallel_apply(self._module_copies[: len(inputs)], inputs, kwargs) output = self.gather(outputs, self.output_device) else: output = self.module.train_step(*inputs, **kwargs) - if torch.is_grad_enabled() and getattr( - self, 'require_backward_grad_sync', True): + if torch.is_grad_enabled() and getattr(self, "require_backward_grad_sync", True): if self.find_unused_parameters: self.reducer.prepare_for_backward(list(_find_tensors(output))) else: self.reducer.prepare_for_backward([]) else: - if ('parrots' not in TORCH_VERSION - and digit_version(TORCH_VERSION) > digit_version('1.2')): + if "parrots" not in TORCH_VERSION and digit_version(TORCH_VERSION) > digit_version("1.2"): self.require_forward_param_sync = False return output @@ -79,34 +75,31 @@ def val_step(self, *inputs, **kwargs): """ # In PyTorch >= 1.7, ``reducer._rebuild_buckets()`` is moved from the # end of backward to the beginning of forward. - if ('parrots' not in TORCH_VERSION - and digit_version(TORCH_VERSION) >= digit_version('1.7') - and self.reducer._rebuild_buckets()): - print_log( - 'Reducer buckets have been rebuilt in this iteration.', - logger='mmcv') - - if getattr(self, 'require_forward_param_sync', True): + if ( + "parrots" not in TORCH_VERSION + and digit_version(TORCH_VERSION) >= digit_version("1.7") + and self.reducer._rebuild_buckets() + ): + print_log("Reducer buckets have been rebuilt in this iteration.", logger="mmcv") + + if getattr(self, "require_forward_param_sync", True): self._sync_params() if self.device_ids: inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids) if len(self.device_ids) == 1: output = self.module.val_step(*inputs[0], **kwargs[0]) else: - outputs = self.parallel_apply( - self._module_copies[:len(inputs)], inputs, kwargs) + outputs = self.parallel_apply(self._module_copies[: len(inputs)], inputs, kwargs) output = self.gather(outputs, self.output_device) else: output = self.module.val_step(*inputs, **kwargs) - if torch.is_grad_enabled() and getattr( - self, 'require_backward_grad_sync', True): + if torch.is_grad_enabled() and getattr(self, "require_backward_grad_sync", True): if self.find_unused_parameters: self.reducer.prepare_for_backward(list(_find_tensors(output))) else: self.reducer.prepare_for_backward([]) else: - if ('parrots' not in TORCH_VERSION - and digit_version(TORCH_VERSION) > digit_version('1.2')): + if "parrots" not in TORCH_VERSION and digit_version(TORCH_VERSION) > digit_version("1.2"): self.require_forward_param_sync = False return output diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/distributed_deprecated.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/distributed_deprecated.py index b039c156..ce87d1f3 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/distributed_deprecated.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/distributed_deprecated.py @@ -2,22 +2,16 @@ import torch import torch.distributed as dist import torch.nn as nn -from torch._utils import (_flatten_dense_tensors, _take_tensors, - _unflatten_dense_tensors) +from torch._utils import _flatten_dense_tensors, _take_tensors, _unflatten_dense_tensors -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, digit_version +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, digit_version from .registry import MODULE_WRAPPERS from .scatter_gather import scatter_kwargs @MODULE_WRAPPERS.register_module() class MMDistributedDataParallel(nn.Module): - - def __init__(self, - module, - dim=0, - broadcast_buffers=True, - bucket_cap_mb=25): + def __init__(self, module, dim=0, broadcast_buffers=True, bucket_cap_mb=25): super(MMDistributedDataParallel, self).__init__() self.module = module self.dim = dim @@ -30,41 +24,34 @@ def _dist_broadcast_coalesced(self, tensors, buffer_size): for tensors in _take_tensors(tensors, buffer_size): flat_tensors = _flatten_dense_tensors(tensors) dist.broadcast(flat_tensors, 0) - for tensor, synced in zip( - tensors, _unflatten_dense_tensors(flat_tensors, tensors)): + for tensor, synced in zip(tensors, _unflatten_dense_tensors(flat_tensors, tensors)): tensor.copy_(synced) def _sync_params(self): module_states = list(self.module.state_dict().values()) if len(module_states) > 0: - self._dist_broadcast_coalesced(module_states, - self.broadcast_bucket_size) + self._dist_broadcast_coalesced(module_states, self.broadcast_bucket_size) if self.broadcast_buffers: - if (TORCH_VERSION != 'parrots' - and digit_version(TORCH_VERSION) < digit_version('1.0')): + if TORCH_VERSION != "parrots" and digit_version(TORCH_VERSION) < digit_version("1.0"): buffers = [b.data for b in self.module._all_buffers()] else: buffers = [b.data for b in self.module.buffers()] if len(buffers) > 0: - self._dist_broadcast_coalesced(buffers, - self.broadcast_bucket_size) + self._dist_broadcast_coalesced(buffers, self.broadcast_bucket_size) def scatter(self, inputs, kwargs, device_ids): return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim) def forward(self, *inputs, **kwargs): - inputs, kwargs = self.scatter(inputs, kwargs, - [torch.cuda.current_device()]) + inputs, kwargs = self.scatter(inputs, kwargs, [torch.cuda.current_device()]) return self.module(*inputs[0], **kwargs[0]) def train_step(self, *inputs, **kwargs): - inputs, kwargs = self.scatter(inputs, kwargs, - [torch.cuda.current_device()]) + inputs, kwargs = self.scatter(inputs, kwargs, [torch.cuda.current_device()]) output = self.module.train_step(*inputs[0], **kwargs[0]) return output def val_step(self, *inputs, **kwargs): - inputs, kwargs = self.scatter(inputs, kwargs, - [torch.cuda.current_device()]) + inputs, kwargs = self.scatter(inputs, kwargs, [torch.cuda.current_device()]) output = self.module.val_step(*inputs[0], **kwargs[0]) return output diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/registry.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/registry.py index a6013c95..ec9e1a8b 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/registry.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/parallel/registry.py @@ -2,7 +2,7 @@ import sys from torch.nn.parallel import DataParallel, DistributedDataParallel -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry MODULE_WRAPPERS = Registry("module wrapper", sys.modules[__name__]) MODULE_WRAPPERS.register_module(module=DataParallel) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/base_module.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/base_module.py index 1ad951b7..e974a048 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/base_module.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/base_module.py @@ -7,8 +7,12 @@ import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.runner.dist_utils import master_only -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.logging import get_logger, logger_initialized, print_log +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner.dist_utils import master_only +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.logging import ( + get_logger, + logger_initialized, + print_log, +) class BaseModule(nn.Module, metaclass=ABCMeta): @@ -58,7 +62,7 @@ def init_weights(self): is_top_level_module = False # check if it is top-level module - if not hasattr(self, '_params_init_info'): + if not hasattr(self, "_params_init_info"): # The `_params_init_info` is used to record the initialization # information of the parameters # the key should be the obj:`nn.Parameter` of model and the value @@ -76,12 +80,12 @@ def init_weights(self): # the corresponding parameter is changed, update related # initialization information for name, param in self.named_parameters(): - self._params_init_info[param][ - 'init_info'] = f'The value is the same before and ' \ - f'after calling `init_weights` ' \ - f'of {self.__class__.__name__} ' - self._params_init_info[param][ - 'tmp_mean_value'] = param.data.mean() + self._params_init_info[param]["init_info"] = ( + f"The value is the same before and " + f"after calling `init_weights` " + f"of {self.__class__.__name__} " + ) + self._params_init_info[param]["tmp_mean_value"] = param.data.mean() # pass `params_init_info` to all submodules # All submodules share the same `params_init_info`, @@ -93,39 +97,35 @@ def init_weights(self): # Get the initialized logger, if not exist, # create a logger named `mmcv` logger_names = list(logger_initialized.keys()) - logger_name = logger_names[0] if logger_names else 'mmcv' + logger_name = logger_names[0] if logger_names else "mmcv" from ..cnn import initialize from ..cnn.utils.weight_init import update_init_info + module_name = self.__class__.__name__ if not self._is_init: if self.init_cfg: - print_log( - f'initialize {module_name} with init_cfg {self.init_cfg}', - logger=logger_name) + print_log(f"initialize {module_name} with init_cfg {self.init_cfg}", logger=logger_name) initialize(self, self.init_cfg) if isinstance(self.init_cfg, dict): # prevent the parameters of # the pre-trained model # from being overwritten by # the `init_weights` - if self.init_cfg['type'] == 'Pretrained': + if self.init_cfg["type"] == "Pretrained": return for m in self.children(): - if hasattr(m, 'init_weights'): + if hasattr(m, "init_weights"): m.init_weights() # users may overload the `init_weights` update_init_info( - m, - init_info=f'Initialized by ' - f'user-defined `init_weights`' - f' in {m.__class__.__name__} ') + m, init_info=f"Initialized by " f"user-defined `init_weights`" f" in {m.__class__.__name__} " + ) self._is_init = True else: - warnings.warn(f'init_weights of {self.__class__.__name__} has ' - f'been called more than once.') + warnings.warn(f"init_weights of {self.__class__.__name__} has " f"been called more than once.") if is_top_level_module: self._dump_init_info(logger_name) @@ -148,25 +148,24 @@ def _dump_init_info(self, logger_name): # dump the information to the logger file if there is a `FileHandler` for handler in logger.handlers: if isinstance(handler, FileHandler): - handler.stream.write( - 'Name of parameter - Initialization information\n') + handler.stream.write("Name of parameter - Initialization information\n") for name, param in self.named_parameters(): handler.stream.write( - f'\n{name} - {param.shape}: ' - f"\n{self._params_init_info[param]['init_info']} \n") + f"\n{name} - {param.shape}: " f"\n{self._params_init_info[param]['init_info']} \n" + ) handler.stream.flush() with_file_handler = True if not with_file_handler: for name, param in self.named_parameters(): print_log( - f'\n{name} - {param.shape}: ' - f"\n{self._params_init_info[param]['init_info']} \n ", - logger=logger_name) + f"\n{name} - {param.shape}: " f"\n{self._params_init_info[param]['init_info']} \n ", + logger=logger_name, + ) def __repr__(self): s = super().__repr__() if self.init_cfg: - s += f'\ninit_cfg={self.init_cfg}' + s += f"\ninit_cfg={self.init_cfg}" return s diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/base_runner.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/base_runner.py index 1d68af1d..6d7c7c4a 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/base_runner.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/base_runner.py @@ -8,7 +8,7 @@ import torch from torch.optim import Optimizer -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv from ..parallel import is_module_wrapper from .checkpoint import load_checkpoint from .dist_utils import get_dist_info @@ -48,55 +48,56 @@ class BaseRunner(metaclass=ABCMeta): max_iters (int, optional): Total training iterations. """ - def __init__(self, - model, - batch_processor=None, - optimizer=None, - work_dir=None, - logger=None, - meta=None, - max_iters=None, - max_epochs=None): + def __init__( + self, + model, + batch_processor=None, + optimizer=None, + work_dir=None, + logger=None, + meta=None, + max_iters=None, + max_epochs=None, + ): if batch_processor is not None: if not callable(batch_processor): - raise TypeError('batch_processor must be callable, ' - f'but got {type(batch_processor)}') - warnings.warn('batch_processor is deprecated, please implement ' - 'train_step() and val_step() in the model instead.') + raise TypeError("batch_processor must be callable, " f"but got {type(batch_processor)}") + warnings.warn( + "batch_processor is deprecated, please implement " "train_step() and val_step() in the model instead." + ) # raise an error is `batch_processor` is not None and # `model.train_step()` exists. if is_module_wrapper(model): _model = model.module else: _model = model - if hasattr(_model, 'train_step') or hasattr(_model, 'val_step'): + if hasattr(_model, "train_step") or hasattr(_model, "val_step"): raise RuntimeError( - 'batch_processor and model.train_step()/model.val_step() ' - 'cannot be both available.') + "batch_processor and model.train_step()/model.val_step() " "cannot be both available." + ) else: - assert hasattr(model, 'train_step') + assert hasattr(model, "train_step") # check the type of `optimizer` if isinstance(optimizer, dict): for name, optim in optimizer.items(): if not isinstance(optim, Optimizer): raise TypeError( - f'optimizer must be a dict of torch.optim.Optimizers, ' - f'but optimizer["{name}"] is a {type(optim)}') + f"optimizer must be a dict of torch.optim.Optimizers, " + f'but optimizer["{name}"] is a {type(optim)}' + ) elif not isinstance(optimizer, Optimizer) and optimizer is not None: raise TypeError( - f'optimizer must be a torch.optim.Optimizer object ' - f'or dict or None, but got {type(optimizer)}') + f"optimizer must be a torch.optim.Optimizer object " f"or dict or None, but got {type(optimizer)}" + ) # check the type of `logger` if not isinstance(logger, logging.Logger): - raise TypeError(f'logger must be a logging.Logger object, ' - f'but got {type(logger)}') + raise TypeError(f"logger must be a logging.Logger object, " f"but got {type(logger)}") # check the type of `meta` if meta is not None and not isinstance(meta, dict): - raise TypeError( - f'meta must be a dict or None, but got {type(meta)}') + raise TypeError(f"meta must be a dict or None, but got {type(meta)}") self.model = model self.batch_processor = batch_processor @@ -113,7 +114,7 @@ def __init__(self, raise TypeError('"work_dir" must be a str or None') # get model name from the model class - if hasattr(self.model, 'module'): + if hasattr(self.model, "module"): self._model_name = self.model.module.__class__.__name__ else: self._model_name = self.model.__class__.__name__ @@ -127,8 +128,7 @@ def __init__(self, self._inner_iter = 0 if max_epochs is not None and max_iters is not None: - raise ValueError( - 'Only one of `max_epochs` or `max_iters` can be set.') + raise ValueError("Only one of `max_epochs` or `max_iters` can be set.") self._max_epochs = max_epochs self._max_iters = max_iters @@ -194,12 +194,7 @@ def run(self, data_loaders, workflow, **kwargs): pass @abstractmethod - def save_checkpoint(self, - out_dir, - filename_tmpl, - save_optimizer=True, - meta=None, - create_symlink=True): + def save_checkpoint(self, out_dir, filename_tmpl, save_optimizer=True, meta=None, create_symlink=True): pass def current_lr(self): @@ -211,14 +206,13 @@ def current_lr(self): method will return a dict. """ if isinstance(self.optimizer, torch.optim.Optimizer): - lr = [group['lr'] for group in self.optimizer.param_groups] + lr = [group["lr"] for group in self.optimizer.param_groups] elif isinstance(self.optimizer, dict): lr = dict() for name, optim in self.optimizer.items(): - lr[name] = [group['lr'] for group in optim.param_groups] + lr[name] = [group["lr"] for group in optim.param_groups] else: - raise RuntimeError( - 'lr is not applicable because optimizer does not exist.') + raise RuntimeError("lr is not applicable because optimizer does not exist.") return lr def current_momentum(self): @@ -233,17 +227,16 @@ def current_momentum(self): def _get_momentum(optimizer): momentums = [] for group in optimizer.param_groups: - if 'momentum' in group.keys(): - momentums.append(group['momentum']) - elif 'betas' in group.keys(): - momentums.append(group['betas'][0]) + if "momentum" in group.keys(): + momentums.append(group["momentum"]) + elif "betas" in group.keys(): + momentums.append(group["betas"][0]) else: momentums.append(0) return momentums if self.optimizer is None: - raise RuntimeError( - 'momentum is not applicable because optimizer does not exist.') + raise RuntimeError("momentum is not applicable because optimizer does not exist.") elif isinstance(self.optimizer, torch.optim.Optimizer): momentums = _get_momentum(self.optimizer) elif isinstance(self.optimizer, dict): @@ -252,7 +245,7 @@ def _get_momentum(optimizer): momentums[name] = _get_momentum(optim) return momentums - def register_hook(self, hook, priority='NORMAL'): + def register_hook(self, hook, priority="NORMAL"): """Register a hook into the hook list. The hook will be inserted into a priority queue, with the specified @@ -266,7 +259,7 @@ def register_hook(self, hook, priority='NORMAL'): Lower value means higher priority. """ assert isinstance(hook, Hook) - if hasattr(hook, 'priority'): + if hasattr(hook, "priority"): raise ValueError('"priority" is a reserved attribute for hooks') priority = get_priority(priority) hook.priority = priority @@ -292,7 +285,7 @@ def register_hook_from_cfg(self, hook_cfg): 'priority' arguments during initialization. """ hook_cfg = hook_cfg.copy() - priority = hook_cfg.pop('priority', 'NORMAL') + priority = hook_cfg.pop("priority", "NORMAL") hook = mmcv.build_from_cfg(hook_cfg, HOOKS) self.register_hook(hook, priority=priority) @@ -315,7 +308,7 @@ def get_hook_info(self): except ValueError: priority = hook.priority classname = hook.__class__.__name__ - hook_info = f'({priority:<12}) {classname:<35}' + hook_info = f"({priority:<12}) {classname:<35}" for trigger_stage in hook.get_triggered_stages(): stage_hook_map[trigger_stage].append(hook_info) @@ -323,85 +316,64 @@ def get_hook_info(self): for stage in Hook.stages: hook_infos = stage_hook_map[stage] if len(hook_infos) > 0: - info = f'{stage}:\n' - info += '\n'.join(hook_infos) - info += '\n -------------------- ' + info = f"{stage}:\n" + info += "\n".join(hook_infos) + info += "\n -------------------- " stage_hook_infos.append(info) - return '\n'.join(stage_hook_infos) - - def load_checkpoint(self, - filename, - map_location='cpu', - strict=False, - revise_keys=[(r'^module.', '')]): - return load_checkpoint( - self.model, - filename, - map_location, - strict, - self.logger, - revise_keys=revise_keys) - - def resume(self, - checkpoint, - resume_optimizer=True, - map_location='default'): - if map_location == 'default': + return "\n".join(stage_hook_infos) + + def load_checkpoint(self, filename, map_location="cpu", strict=False, revise_keys=[(r"^module.", "")]): + return load_checkpoint(self.model, filename, map_location, strict, self.logger, revise_keys=revise_keys) + + def resume(self, checkpoint, resume_optimizer=True, map_location="default"): + if map_location == "default": if torch.cuda.is_available(): device_id = torch.cuda.current_device() checkpoint = self.load_checkpoint( - checkpoint, - map_location=lambda storage, loc: storage.cuda(device_id)) + checkpoint, map_location=lambda storage, loc: storage.cuda(device_id) + ) else: checkpoint = self.load_checkpoint(checkpoint) else: - checkpoint = self.load_checkpoint( - checkpoint, map_location=map_location) + checkpoint = self.load_checkpoint(checkpoint, map_location=map_location) - self._epoch = checkpoint['meta']['epoch'] - self._iter = checkpoint['meta']['iter'] + self._epoch = checkpoint["meta"]["epoch"] + self._iter = checkpoint["meta"]["iter"] if self.meta is None: self.meta = {} - self.meta.setdefault('hook_msgs', {}) + self.meta.setdefault("hook_msgs", {}) # load `last_ckpt`, `best_score`, `best_ckpt`, etc. for hook messages - self.meta['hook_msgs'].update(checkpoint['meta'].get('hook_msgs', {})) + self.meta["hook_msgs"].update(checkpoint["meta"].get("hook_msgs", {})) # Re-calculate the number of iterations when resuming # models with different number of GPUs - if 'config' in checkpoint['meta']: - config = mmcv.Config.fromstring( - checkpoint['meta']['config'], file_format='.py') - previous_gpu_ids = config.get('gpu_ids', None) - if previous_gpu_ids and len(previous_gpu_ids) > 0 and len( - previous_gpu_ids) != self.world_size: - self._iter = int(self._iter * len(previous_gpu_ids) / - self.world_size) - self.logger.info('the iteration number is changed due to ' - 'change of GPU number') + if "config" in checkpoint["meta"]: + config = mmcv.Config.fromstring(checkpoint["meta"]["config"], file_format=".py") + previous_gpu_ids = config.get("gpu_ids", None) + if previous_gpu_ids and len(previous_gpu_ids) > 0 and len(previous_gpu_ids) != self.world_size: + self._iter = int(self._iter * len(previous_gpu_ids) / self.world_size) + self.logger.info("the iteration number is changed due to " "change of GPU number") # resume meta information meta - self.meta = checkpoint['meta'] + self.meta = checkpoint["meta"] - if 'optimizer' in checkpoint and resume_optimizer: + if "optimizer" in checkpoint and resume_optimizer: if isinstance(self.optimizer, Optimizer): - self.optimizer.load_state_dict(checkpoint['optimizer']) + self.optimizer.load_state_dict(checkpoint["optimizer"]) elif isinstance(self.optimizer, dict): for k in self.optimizer.keys(): - self.optimizer[k].load_state_dict( - checkpoint['optimizer'][k]) + self.optimizer[k].load_state_dict(checkpoint["optimizer"][k]) else: - raise TypeError( - 'Optimizer should be dict or torch.optim.Optimizer ' - f'but got {type(self.optimizer)}') + raise TypeError("Optimizer should be dict or torch.optim.Optimizer " f"but got {type(self.optimizer)}") - self.logger.info('resumed epoch %d, iter %d', self.epoch, self.iter) + self.logger.info("resumed epoch %d, iter %d", self.epoch, self.iter) def register_lr_hook(self, lr_config): if lr_config is None: return elif isinstance(lr_config, dict): - assert 'policy' in lr_config - policy_type = lr_config.pop('policy') + assert "policy" in lr_config + policy_type = lr_config.pop("policy") # If the type of policy is all in lower case, e.g., 'cyclic', # then its first letter will be capitalized, e.g., to be 'Cyclic'. # This is for the convenient usage of Lr updater. @@ -410,19 +382,19 @@ def register_lr_hook(self, lr_config): # the string will not be changed if it contains capital letters. if policy_type == policy_type.lower(): policy_type = policy_type.title() - hook_type = policy_type + 'LrUpdaterHook' - lr_config['type'] = hook_type + hook_type = policy_type + "LrUpdaterHook" + lr_config["type"] = hook_type hook = mmcv.build_from_cfg(lr_config, HOOKS) else: hook = lr_config - self.register_hook(hook, priority='VERY_HIGH') + self.register_hook(hook, priority="VERY_HIGH") def register_momentum_hook(self, momentum_config): if momentum_config is None: return if isinstance(momentum_config, dict): - assert 'policy' in momentum_config - policy_type = momentum_config.pop('policy') + assert "policy" in momentum_config + policy_type = momentum_config.pop("policy") # If the type of policy is all in lower case, e.g., 'cyclic', # then its first letter will be capitalized, e.g., to be 'Cyclic'. # This is for the convenient usage of momentum updater. @@ -431,41 +403,40 @@ def register_momentum_hook(self, momentum_config): # the string will not be changed if it contains capital letters. if policy_type == policy_type.lower(): policy_type = policy_type.title() - hook_type = policy_type + 'MomentumUpdaterHook' - momentum_config['type'] = hook_type + hook_type = policy_type + "MomentumUpdaterHook" + momentum_config["type"] = hook_type hook = mmcv.build_from_cfg(momentum_config, HOOKS) else: hook = momentum_config - self.register_hook(hook, priority='HIGH') + self.register_hook(hook, priority="HIGH") def register_optimizer_hook(self, optimizer_config): if optimizer_config is None: return if isinstance(optimizer_config, dict): - optimizer_config.setdefault('type', 'OptimizerHook') + optimizer_config.setdefault("type", "OptimizerHook") hook = mmcv.build_from_cfg(optimizer_config, HOOKS) else: hook = optimizer_config - self.register_hook(hook, priority='ABOVE_NORMAL') + self.register_hook(hook, priority="ABOVE_NORMAL") def register_checkpoint_hook(self, checkpoint_config): if checkpoint_config is None: return if isinstance(checkpoint_config, dict): - checkpoint_config.setdefault('type', 'CheckpointHook') + checkpoint_config.setdefault("type", "CheckpointHook") hook = mmcv.build_from_cfg(checkpoint_config, HOOKS) else: hook = checkpoint_config - self.register_hook(hook, priority='NORMAL') + self.register_hook(hook, priority="NORMAL") def register_logger_hooks(self, log_config): if log_config is None: return - log_interval = log_config['interval'] - for info in log_config['hooks']: - logger_hook = mmcv.build_from_cfg( - info, HOOKS, default_args=dict(interval=log_interval)) - self.register_hook(logger_hook, priority='VERY_LOW') + log_interval = log_config["interval"] + for info in log_config["hooks"]: + logger_hook = mmcv.build_from_cfg(info, HOOKS, default_args=dict(interval=log_interval)) + self.register_hook(logger_hook, priority="VERY_LOW") def register_timer_hook(self, timer_config): if timer_config is None: @@ -475,7 +446,7 @@ def register_timer_hook(self, timer_config): hook = mmcv.build_from_cfg(timer_config_, HOOKS) else: hook = timer_config - self.register_hook(hook, priority='LOW') + self.register_hook(hook, priority="LOW") def register_custom_hooks(self, custom_config): if custom_config is None: @@ -488,26 +459,28 @@ def register_custom_hooks(self, custom_config): if isinstance(item, dict): self.register_hook_from_cfg(item) else: - self.register_hook(item, priority='NORMAL') + self.register_hook(item, priority="NORMAL") def register_profiler_hook(self, profiler_config): if profiler_config is None: return if isinstance(profiler_config, dict): - profiler_config.setdefault('type', 'ProfilerHook') + profiler_config.setdefault("type", "ProfilerHook") hook = mmcv.build_from_cfg(profiler_config, HOOKS) else: hook = profiler_config self.register_hook(hook) - def register_training_hooks(self, - lr_config, - optimizer_config=None, - checkpoint_config=None, - log_config=None, - momentum_config=None, - timer_config=dict(type='IterTimerHook'), - custom_hooks_config=None): + def register_training_hooks( + self, + lr_config, + optimizer_config=None, + checkpoint_config=None, + log_config=None, + momentum_config=None, + timer_config=dict(type="IterTimerHook"), + custom_hooks_config=None, + ): """Register default and custom hooks for training. Default and custom hooks include: diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/checkpoint.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/checkpoint.py index 016ac125..a9d66915 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/checkpoint.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/checkpoint.py @@ -15,24 +15,22 @@ from torch.optim import Optimizer from torch.utils import model_zoo -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv from ..fileio import FileClient from ..fileio import load as load_file from ..parallel import is_module_wrapper from ..utils import mkdir_or_exist from .dist_utils import get_dist_info -ENV_MMCV_HOME = 'MMCV_HOME' -ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME' -DEFAULT_CACHE_DIR = '~/.cache' +ENV_MMCV_HOME = "MMCV_HOME" +ENV_XDG_CACHE_HOME = "XDG_CACHE_HOME" +DEFAULT_CACHE_DIR = "~/.cache" def _get_mmcv_home(): mmcv_home = os.path.expanduser( - os.getenv( - ENV_MMCV_HOME, - os.path.join( - os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'mmcv'))) + os.getenv(ENV_MMCV_HOME, os.path.join(os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), "mmcv")) + ) mkdir_or_exist(mmcv_home) return mmcv_home @@ -58,46 +56,40 @@ def load_state_dict(module, state_dict, strict=False, logger=None): all_missing_keys = [] err_msg = [] - metadata = getattr(state_dict, '_metadata', None) + metadata = getattr(state_dict, "_metadata", None) state_dict = state_dict.copy() if metadata is not None: state_dict._metadata = metadata # use _load_from_state_dict to enable checkpoint version control - def load(module, prefix=''): + def load(module, prefix=""): # recursively check parallel module in case that the model has a # complicated structure, e.g., nn.Module(nn.Module(DDP)) if is_module_wrapper(module): module = module.module - local_metadata = {} if metadata is None else metadata.get( - prefix[:-1], {}) - module._load_from_state_dict(state_dict, prefix, local_metadata, True, - all_missing_keys, unexpected_keys, - err_msg) + local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) + module._load_from_state_dict( + state_dict, prefix, local_metadata, True, all_missing_keys, unexpected_keys, err_msg + ) for name, child in module._modules.items(): if child is not None: - load(child, prefix + name + '.') + load(child, prefix + name + ".") load(module) load = None # break load->load reference cycle # ignore "num_batches_tracked" of BN layers - missing_keys = [ - key for key in all_missing_keys if 'num_batches_tracked' not in key - ] + missing_keys = [key for key in all_missing_keys if "num_batches_tracked" not in key] if unexpected_keys: - err_msg.append('unexpected key in source ' - f'state_dict: {", ".join(unexpected_keys)}\n') + err_msg.append("unexpected key in source " f'state_dict: {", ".join(unexpected_keys)}\n') if missing_keys: - err_msg.append( - f'missing keys in source state_dict: {", ".join(missing_keys)}\n') + err_msg.append(f'missing keys in source state_dict: {", ".join(missing_keys)}\n') rank, _ = get_dist_info() if len(err_msg) > 0 and rank == 0: - err_msg.insert( - 0, 'The model and loaded state dict do not match exactly\n') - err_msg = '\n'.join(err_msg) + err_msg.insert(0, "The model and loaded state dict do not match exactly\n") + err_msg = "\n".join(err_msg) if strict: raise RuntimeError(err_msg) elif logger is not None: @@ -111,19 +103,19 @@ def get_torchvision_models(): for _, name, ispkg in pkgutil.walk_packages(torchvision.models.__path__): if ispkg: continue - _zoo = import_module(f'torchvision.models.{name}') - if hasattr(_zoo, 'model_urls'): - _urls = getattr(_zoo, 'model_urls') + _zoo = import_module(f"torchvision.models.{name}") + if hasattr(_zoo, "model_urls"): + _urls = getattr(_zoo, "model_urls") model_urls.update(_urls) return model_urls def get_external_models(): mmcv_home = _get_mmcv_home() - default_json_path = osp.join(mmcv.__path__[0], 'model_zoo/open_mmlab.json') + default_json_path = osp.join(mmcv.__path__[0], "model_zoo/open_mmlab.json") default_urls = load_file(default_json_path) assert isinstance(default_urls, dict) - external_json_path = osp.join(mmcv_home, 'open_mmlab.json') + external_json_path = osp.join(mmcv_home, "open_mmlab.json") if osp.exists(external_json_path): external_urls = load_file(external_json_path) assert isinstance(external_urls, dict) @@ -133,15 +125,14 @@ def get_external_models(): def get_mmcls_models(): - mmcls_json_path = osp.join(mmcv.__path__[0], 'model_zoo/mmcls.json') + mmcls_json_path = osp.join(mmcv.__path__[0], "model_zoo/mmcls.json") mmcls_urls = load_file(mmcls_json_path) return mmcls_urls def get_deprecated_model_names(): - deprecate_json_path = osp.join(mmcv.__path__[0], - 'model_zoo/deprecated.json') + deprecate_json_path = osp.join(mmcv.__path__[0], "model_zoo/deprecated.json") deprecate_urls = load_file(deprecate_json_path) assert isinstance(deprecate_urls, dict) @@ -149,10 +140,10 @@ def get_deprecated_model_names(): def _process_mmcls_checkpoint(checkpoint): - state_dict = checkpoint['state_dict'] + state_dict = checkpoint["state_dict"] new_state_dict = OrderedDict() for k, v in state_dict.items(): - if k.startswith('backbone.'): + if k.startswith("backbone."): new_state_dict[k[9:]] = v new_checkpoint = dict(state_dict=new_state_dict) @@ -175,11 +166,11 @@ def _register_scheme(cls, prefixes, loader, force=False): cls._schemes[prefix] = loader else: raise KeyError( - f'{prefix} is already registered as a loader backend, ' - 'add "force=True" if you want to override it') + f"{prefix} is already registered as a loader backend, " + 'add "force=True" if you want to override it' + ) # sort, longer prefixes take priority - cls._schemes = OrderedDict( - sorted(cls._schemes.items(), key=lambda t: t[0], reverse=True)) + cls._schemes = OrderedDict(sorted(cls._schemes.items(), key=lambda t: t[0], reverse=True)) @classmethod def register_scheme(cls, prefixes, loader=None, force=False): @@ -240,12 +231,11 @@ def load_checkpoint(cls, filename, map_location=None, logger=None): checkpoint_loader = cls._get_checkpoint_loader(filename) class_name = checkpoint_loader.__name__ - mmcv.print_log( - f'load checkpoint from {class_name[10:]} path: {filename}', logger) + mmcv.print_log(f"load checkpoint from {class_name[10:]} path: {filename}", logger) return checkpoint_loader(filename, map_location) -@CheckpointLoader.register_scheme(prefixes='') +@CheckpointLoader.register_scheme(prefixes="") def load_from_local(filename, map_location): """load checkpoint by local file path. @@ -258,12 +248,12 @@ def load_from_local(filename, map_location): """ if not osp.isfile(filename): - raise IOError(f'{filename} is not a checkpoint file') + raise IOError(f"{filename} is not a checkpoint file") checkpoint = torch.load(filename, map_location=map_location) return checkpoint -@CheckpointLoader.register_scheme(prefixes=('http://', 'https://')) +@CheckpointLoader.register_scheme(prefixes=("http://", "https://")) def load_from_http(filename, map_location=None, model_dir=None): """load checkpoint through HTTP or HTTPS scheme path. In distributed setting, this function only download checkpoint at local rank 0. @@ -279,19 +269,17 @@ def load_from_http(filename, map_location=None, model_dir=None): dict or OrderedDict: The loaded checkpoint. """ rank, world_size = get_dist_info() - rank = int(os.environ.get('LOCAL_RANK', rank)) + rank = int(os.environ.get("LOCAL_RANK", rank)) if rank == 0: - checkpoint = model_zoo.load_url( - filename, model_dir=model_dir, map_location=map_location) + checkpoint = model_zoo.load_url(filename, model_dir=model_dir, map_location=map_location) if world_size > 1: torch.distributed.barrier() if rank > 0: - checkpoint = model_zoo.load_url( - filename, model_dir=model_dir, map_location=map_location) + checkpoint = model_zoo.load_url(filename, model_dir=model_dir, map_location=map_location) return checkpoint -@CheckpointLoader.register_scheme(prefixes='pavi://') +@CheckpointLoader.register_scheme(prefixes="pavi://") def load_from_pavi(filename, map_location=None): """load checkpoint through the file path prefixed with pavi. In distributed setting, this function download ckpt at all ranks to different temporary @@ -305,15 +293,13 @@ def load_from_pavi(filename, map_location=None): Returns: dict or OrderedDict: The loaded checkpoint. """ - assert filename.startswith('pavi://'), \ - f'Expected filename startswith `pavi://`, but get {filename}' + assert filename.startswith("pavi://"), f"Expected filename startswith `pavi://`, but get {filename}" model_path = filename[7:] try: from pavi import modelcloud except ImportError: - raise ImportError( - 'Please install pavi to load checkpoint from modelcloud.') + raise ImportError("Please install pavi to load checkpoint from modelcloud.") model = modelcloud.get(model_path) with TemporaryDirectory() as tmp_dir: @@ -323,8 +309,8 @@ def load_from_pavi(filename, map_location=None): return checkpoint -@CheckpointLoader.register_scheme(prefixes='s3://') -def load_from_ceph(filename, map_location=None, backend='petrel'): +@CheckpointLoader.register_scheme(prefixes="s3://") +def load_from_ceph(filename, map_location=None, backend="petrel"): """load checkpoint through the file path prefixed with s3. In distributed setting, this function download ckpt at all ranks to different temporary directories. @@ -342,13 +328,12 @@ def load_from_ceph(filename, map_location=None, backend='petrel'): Returns: dict or OrderedDict: The loaded checkpoint. """ - allowed_backends = ['ceph', 'petrel'] + allowed_backends = ["ceph", "petrel"] if backend not in allowed_backends: - raise ValueError(f'Load from Backend {backend} is not supported.') + raise ValueError(f"Load from Backend {backend} is not supported.") - if backend == 'ceph': - warnings.warn( - 'CephBackend will be deprecated, please use PetrelBackend instead') + if backend == "ceph": + warnings.warn("CephBackend will be deprecated, please use PetrelBackend instead") # CephClient and PetrelBackend have the same prefix 's3://' and the latter # will be chosen as default. If PetrelBackend can not be instantiated @@ -364,7 +349,7 @@ def load_from_ceph(filename, map_location=None, backend='petrel'): return checkpoint -@CheckpointLoader.register_scheme(prefixes=('modelzoo://', 'torchvision://')) +@CheckpointLoader.register_scheme(prefixes=("modelzoo://", "torchvision://")) def load_from_torchvision(filename, map_location=None): """load checkpoint through the file path prefixed with modelzoo or torchvision. @@ -378,16 +363,15 @@ def load_from_torchvision(filename, map_location=None): dict or OrderedDict: The loaded checkpoint. """ model_urls = get_torchvision_models() - if filename.startswith('modelzoo://'): - warnings.warn('The URL scheme of "modelzoo://" is deprecated, please ' - 'use "torchvision://" instead') + if filename.startswith("modelzoo://"): + warnings.warn('The URL scheme of "modelzoo://" is deprecated, please ' 'use "torchvision://" instead') model_name = filename[11:] else: model_name = filename[14:] return load_from_http(model_urls[model_name], map_location=map_location) -@CheckpointLoader.register_scheme(prefixes=('open-mmlab://', 'openmmlab://')) +@CheckpointLoader.register_scheme(prefixes=("open-mmlab://", "openmmlab://")) def load_from_openmmlab(filename, map_location=None): """load checkpoint through the file path prefixed with open-mmlab or openmmlab. @@ -403,31 +387,32 @@ def load_from_openmmlab(filename, map_location=None): """ model_urls = get_external_models() - prefix_str = 'open-mmlab://' + prefix_str = "open-mmlab://" if filename.startswith(prefix_str): model_name = filename[13:] else: model_name = filename[12:] - prefix_str = 'openmmlab://' + prefix_str = "openmmlab://" deprecated_urls = get_deprecated_model_names() if model_name in deprecated_urls: - warnings.warn(f'{prefix_str}{model_name} is deprecated in favor ' - f'of {prefix_str}{deprecated_urls[model_name]}') + warnings.warn( + f"{prefix_str}{model_name} is deprecated in favor " f"of {prefix_str}{deprecated_urls[model_name]}" + ) model_name = deprecated_urls[model_name] model_url = model_urls[model_name] # check if is url - if model_url.startswith(('http://', 'https://')): + if model_url.startswith(("http://", "https://")): checkpoint = load_from_http(model_url, map_location=map_location) else: filename = osp.join(_get_mmcv_home(), model_url) if not osp.isfile(filename): - raise IOError(f'{filename} is not a checkpoint file') + raise IOError(f"{filename} is not a checkpoint file") checkpoint = torch.load(filename, map_location=map_location) return checkpoint -@CheckpointLoader.register_scheme(prefixes='mmcls://') +@CheckpointLoader.register_scheme(prefixes="mmcls://") def load_from_mmcls(filename, map_location=None): """load checkpoint through the file path prefixed with mmcls. @@ -441,8 +426,7 @@ def load_from_mmcls(filename, map_location=None): model_urls = get_mmcls_models() model_name = filename[8:] - checkpoint = load_from_http( - model_urls[model_name], map_location=map_location) + checkpoint = load_from_http(model_urls[model_name], map_location=map_location) checkpoint = _process_mmcls_checkpoint(checkpoint) return checkpoint @@ -483,29 +467,21 @@ def _load_checkpoint_with_prefix(prefix, filename, map_location=None): checkpoint = _load_checkpoint(filename, map_location=map_location) - if 'state_dict' in checkpoint: - state_dict = checkpoint['state_dict'] + if "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] else: state_dict = checkpoint - if not prefix.endswith('.'): - prefix += '.' + if not prefix.endswith("."): + prefix += "." prefix_len = len(prefix) - state_dict = { - k[prefix_len:]: v - for k, v in state_dict.items() if k.startswith(prefix) - } + state_dict = {k[prefix_len:]: v for k, v in state_dict.items() if k.startswith(prefix)} - assert state_dict, f'{prefix} is not in the pretrained model' + assert state_dict, f"{prefix} is not in the pretrained model" return state_dict -def load_checkpoint(model, - filename, - map_location=None, - strict=False, - logger=None, - revise_keys=[(r'^module\.', '')]): +def load_checkpoint(model, filename, map_location=None, strict=False, logger=None, revise_keys=[(r"^module\.", "")]): """Load checkpoint from a file or URI. Args: @@ -528,20 +504,17 @@ def load_checkpoint(model, checkpoint = _load_checkpoint(filename, map_location, logger) # OrderedDict is a subclass of dict if not isinstance(checkpoint, dict): - raise RuntimeError( - f'No state_dict found in checkpoint file {filename}') + raise RuntimeError(f"No state_dict found in checkpoint file {filename}") # get state_dict from checkpoint - if 'state_dict' in checkpoint: - state_dict = checkpoint['state_dict'] + if "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] else: state_dict = checkpoint # strip prefix of state_dict - metadata = getattr(state_dict, '_metadata', OrderedDict()) + metadata = getattr(state_dict, "_metadata", OrderedDict()) for p, r in revise_keys: - state_dict = OrderedDict( - {re.sub(p, r, k): v - for k, v in state_dict.items()}) + state_dict = OrderedDict({re.sub(p, r, k): v for k, v in state_dict.items()}) # Keep metadata in state_dict state_dict._metadata = metadata @@ -563,7 +536,7 @@ def weights_to_cpu(state_dict): for key, val in state_dict.items(): state_dict_cpu[key] = val.cpu() # Keep metadata in state_dict - state_dict_cpu._metadata = getattr(state_dict, '_metadata', OrderedDict()) + state_dict_cpu._metadata = getattr(state_dict, "_metadata", OrderedDict()) return state_dict_cpu @@ -587,7 +560,7 @@ def _save_to_state_dict(module, destination, prefix, keep_vars): destination[prefix + name] = buf if keep_vars else buf.detach() -def get_state_dict(module, destination=None, prefix='', keep_vars=False): +def get_state_dict(module, destination=None, prefix="", keep_vars=False): """Returns a dictionary containing a whole state of the module. Both parameters and persistent buffers (e.g. running averages) are @@ -617,13 +590,11 @@ def get_state_dict(module, destination=None, prefix='', keep_vars=False): if destination is None: destination = OrderedDict() destination._metadata = OrderedDict() - destination._metadata[prefix[:-1]] = local_metadata = dict( - version=module._version) + destination._metadata[prefix[:-1]] = local_metadata = dict(version=module._version) _save_to_state_dict(module, destination, prefix, keep_vars) for name, child in module._modules.items(): if child is not None: - get_state_dict( - child, destination, prefix + name + '.', keep_vars=keep_vars) + get_state_dict(child, destination, prefix + name + ".", keep_vars=keep_vars) for hook in module._state_dict_hooks.values(): hook_result = hook(module, destination, prefix, local_metadata) if hook_result is not None: @@ -631,11 +602,7 @@ def get_state_dict(module, destination=None, prefix='', keep_vars=False): return destination -def save_checkpoint(model, - filename, - optimizer=None, - meta=None, - file_client_args=None): +def save_checkpoint(model, filename, optimizer=None, meta=None, file_client_args=None): """Save checkpoint to file. The checkpoint will have 3 fields: ``meta``, ``state_dict`` and @@ -654,39 +621,35 @@ def save_checkpoint(model, if meta is None: meta = {} elif not isinstance(meta, dict): - raise TypeError(f'meta must be a dict or None, but got {type(meta)}') + raise TypeError(f"meta must be a dict or None, but got {type(meta)}") meta.update(mmcv_version=mmcv.__version__, time=time.asctime()) if is_module_wrapper(model): model = model.module - if hasattr(model, 'CLASSES') and model.CLASSES is not None: + if hasattr(model, "CLASSES") and model.CLASSES is not None: # save class name to the meta meta.update(CLASSES=model.CLASSES) - checkpoint = { - 'meta': meta, - 'state_dict': weights_to_cpu(get_state_dict(model)) - } + checkpoint = {"meta": meta, "state_dict": weights_to_cpu(get_state_dict(model))} # save optimizer state dict in the checkpoint if isinstance(optimizer, Optimizer): - checkpoint['optimizer'] = optimizer.state_dict() + checkpoint["optimizer"] = optimizer.state_dict() elif isinstance(optimizer, dict): - checkpoint['optimizer'] = {} + checkpoint["optimizer"] = {} for name, optim in optimizer.items(): - checkpoint['optimizer'][name] = optim.state_dict() + checkpoint["optimizer"][name] = optim.state_dict() - if filename.startswith('pavi://'): + if filename.startswith("pavi://"): if file_client_args is not None: raise ValueError( - 'file_client_args should be "None" if filename starts with' - f'"pavi://", but got {file_client_args}') + 'file_client_args should be "None" if filename starts with' f'"pavi://", but got {file_client_args}' + ) try: from pavi import modelcloud from pavi import exception except ImportError: - raise ImportError( - 'Please install pavi to load checkpoint from modelcloud.') + raise ImportError("Please install pavi to load checkpoint from modelcloud.") model_path = filename[7:] root = modelcloud.Folder() model_dir, model_name = osp.split(model_path) @@ -696,7 +659,7 @@ def save_checkpoint(model, model = root.create_training_model(model_dir) with TemporaryDirectory() as tmp_dir: checkpoint_file = osp.join(tmp_dir, model_name) - with open(checkpoint_file, 'wb') as f: + with open(checkpoint_file, "wb") as f: torch.save(checkpoint, f) f.flush() model.create_file(checkpoint_file, name=model_name) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/default_constructor.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/default_constructor.py index f82fa028..cad735f4 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/default_constructor.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/default_constructor.py @@ -9,7 +9,7 @@ class DefaultRunnerConstructor: For example, We can inject some new properties and functions for `Runner`. Example: - >>> from comfy_controlnet_preprocessors.uniformer.mmcv.runner import RUNNER_BUILDERS, build_runner + >>> from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import RUNNER_BUILDERS, build_runner >>> # Define a new RunnerReconstructor >>> @RUNNER_BUILDERS.register_module() >>> class MyRunnerConstructor: @@ -35,8 +35,7 @@ class DefaultRunnerConstructor: def __init__(self, runner_cfg, default_args=None): if not isinstance(runner_cfg, dict): - raise TypeError('runner_cfg should be a dict', - f'but got {type(runner_cfg)}') + raise TypeError("runner_cfg should be a dict", f"but got {type(runner_cfg)}") self.runner_cfg = runner_cfg self.default_args = default_args diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/epoch_based_runner.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/epoch_based_runner.py index bee059de..f2b6a796 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/epoch_based_runner.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/epoch_based_runner.py @@ -7,7 +7,7 @@ import torch -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv from .base_runner import BaseRunner from .builder import RUNNERS from .checkpoint import save_checkpoint @@ -23,51 +23,48 @@ class EpochBasedRunner(BaseRunner): def run_iter(self, data_batch, train_mode, **kwargs): if self.batch_processor is not None: - outputs = self.batch_processor( - self.model, data_batch, train_mode=train_mode, **kwargs) + outputs = self.batch_processor(self.model, data_batch, train_mode=train_mode, **kwargs) elif train_mode: - outputs = self.model.train_step(data_batch, self.optimizer, - **kwargs) + outputs = self.model.train_step(data_batch, self.optimizer, **kwargs) else: outputs = self.model.val_step(data_batch, self.optimizer, **kwargs) if not isinstance(outputs, dict): - raise TypeError('"batch_processor()" or "model.train_step()"' - 'and "model.val_step()" must return a dict') - if 'log_vars' in outputs: - self.log_buffer.update(outputs['log_vars'], outputs['num_samples']) + raise TypeError('"batch_processor()" or "model.train_step()"' 'and "model.val_step()" must return a dict') + if "log_vars" in outputs: + self.log_buffer.update(outputs["log_vars"], outputs["num_samples"]) self.outputs = outputs def train(self, data_loader, **kwargs): self.model.train() - self.mode = 'train' + self.mode = "train" self.data_loader = data_loader self._max_iters = self._max_epochs * len(self.data_loader) - self.call_hook('before_train_epoch') + self.call_hook("before_train_epoch") time.sleep(2) # Prevent possible deadlock during epoch transition for i, data_batch in enumerate(self.data_loader): self._inner_iter = i - self.call_hook('before_train_iter') + self.call_hook("before_train_iter") self.run_iter(data_batch, train_mode=True, **kwargs) - self.call_hook('after_train_iter') + self.call_hook("after_train_iter") self._iter += 1 - self.call_hook('after_train_epoch') + self.call_hook("after_train_epoch") self._epoch += 1 @torch.no_grad() def val(self, data_loader, **kwargs): self.model.eval() - self.mode = 'val' + self.mode = "val" self.data_loader = data_loader - self.call_hook('before_val_epoch') + self.call_hook("before_val_epoch") time.sleep(2) # Prevent possible deadlock during epoch transition for i, data_batch in enumerate(self.data_loader): self._inner_iter = i - self.call_hook('before_val_iter') + self.call_hook("before_val_iter") self.run_iter(data_batch, train_mode=False) - self.call_hook('after_val_iter') + self.call_hook("after_val_iter") - self.call_hook('after_val_epoch') + self.call_hook("after_val_epoch") def run(self, data_loaders, workflow, max_epochs=None, **kwargs): """Start running. @@ -85,56 +82,46 @@ def run(self, data_loaders, workflow, max_epochs=None, **kwargs): assert len(data_loaders) == len(workflow) if max_epochs is not None: warnings.warn( - 'setting max_epochs in run is deprecated, ' - 'please set max_epochs in runner_config', DeprecationWarning) + "setting max_epochs in run is deprecated, " "please set max_epochs in runner_config", + DeprecationWarning, + ) self._max_epochs = max_epochs - assert self._max_epochs is not None, ( - 'max_epochs must be specified during instantiation') + assert self._max_epochs is not None, "max_epochs must be specified during instantiation" for i, flow in enumerate(workflow): mode, epochs = flow - if mode == 'train': + if mode == "train": self._max_iters = self._max_epochs * len(data_loaders[i]) break - work_dir = self.work_dir if self.work_dir is not None else 'NONE' - self.logger.info('Start running, host: %s, work_dir: %s', - get_host_info(), work_dir) - self.logger.info('Hooks will be executed in the following order:\n%s', - self.get_hook_info()) - self.logger.info('workflow: %s, max: %d epochs', workflow, - self._max_epochs) - self.call_hook('before_run') + work_dir = self.work_dir if self.work_dir is not None else "NONE" + self.logger.info("Start running, host: %s, work_dir: %s", get_host_info(), work_dir) + self.logger.info("Hooks will be executed in the following order:\n%s", self.get_hook_info()) + self.logger.info("workflow: %s, max: %d epochs", workflow, self._max_epochs) + self.call_hook("before_run") while self.epoch < self._max_epochs: for i, flow in enumerate(workflow): mode, epochs = flow if isinstance(mode, str): # self.train() if not hasattr(self, mode): - raise ValueError( - f'runner has no method named "{mode}" to run an ' - 'epoch') + raise ValueError(f'runner has no method named "{mode}" to run an ' "epoch") epoch_runner = getattr(self, mode) else: - raise TypeError( - 'mode in workflow must be a str, but got {}'.format( - type(mode))) + raise TypeError("mode in workflow must be a str, but got {}".format(type(mode))) for _ in range(epochs): - if mode == 'train' and self.epoch >= self._max_epochs: + if mode == "train" and self.epoch >= self._max_epochs: break epoch_runner(data_loaders[i], **kwargs) time.sleep(1) # wait for some hooks like loggers to finish - self.call_hook('after_run') - - def save_checkpoint(self, - out_dir, - filename_tmpl='epoch_{}.pth', - save_optimizer=True, - meta=None, - create_symlink=True): + self.call_hook("after_run") + + def save_checkpoint( + self, out_dir, filename_tmpl="epoch_{}.pth", save_optimizer=True, meta=None, create_symlink=True + ): """Save the checkpoint. Args: @@ -153,8 +140,7 @@ def save_checkpoint(self, if meta is None: meta = {} elif not isinstance(meta, dict): - raise TypeError( - f'meta should be a dict or None, but got {type(meta)}') + raise TypeError(f"meta should be a dict or None, but got {type(meta)}") if self.meta is not None: meta.update(self.meta) # Note: meta.update(self.meta) should be done before @@ -170,8 +156,8 @@ def save_checkpoint(self, # in some environments, `os.symlink` is not supported, you may need to # set `create_symlink` to False if create_symlink: - dst_file = osp.join(out_dir, 'latest.pth') - if platform.system() != 'Windows': + dst_file = osp.join(out_dir, "latest.pth") + if platform.system() != "Windows": mmcv.symlink(filename, dst_file) else: shutil.copy(filepath, dst_file) @@ -182,6 +168,5 @@ class Runner(EpochBasedRunner): """Deprecated name of EpochBasedRunner.""" def __init__(self, *args, **kwargs): - warnings.warn( - 'Runner was deprecated, please use EpochBasedRunner instead') + warnings.warn("Runner was deprecated, please use EpochBasedRunner instead") super().__init__(*args, **kwargs) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/fp16_utils.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/fp16_utils.py index da9b364a..36676c0d 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/fp16_utils.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/fp16_utils.py @@ -8,7 +8,7 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, digit_version +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, digit_version from .dist_utils import allreduce_grads as _allreduce_grads try: @@ -41,13 +41,9 @@ def cast_tensor_type(inputs, src_type, dst_type): elif isinstance(inputs, np.ndarray): return inputs elif isinstance(inputs, abc.Mapping): - return type(inputs)({ - k: cast_tensor_type(v, src_type, dst_type) - for k, v in inputs.items() - }) + return type(inputs)({k: cast_tensor_type(v, src_type, dst_type) for k, v in inputs.items()}) elif isinstance(inputs, abc.Iterable): - return type(inputs)( - cast_tensor_type(item, src_type, dst_type) for item in inputs) + return type(inputs)(cast_tensor_type(item, src_type, dst_type) for item in inputs) else: return inputs @@ -86,15 +82,13 @@ def auto_fp16(apply_to=None, out_fp32=False): """ def auto_fp16_wrapper(old_func): - @functools.wraps(old_func) def new_func(*args, **kwargs): # check if the module has set the attribute `fp16_enabled`, if not, # just fallback to the original method. if not isinstance(args[0], torch.nn.Module): - raise TypeError('@auto_fp16 can only be used to decorate the ' - 'method of nn.Module') - if not (hasattr(args[0], 'fp16_enabled') and args[0].fp16_enabled): + raise TypeError("@auto_fp16 can only be used to decorate the " "method of nn.Module") + if not (hasattr(args[0], "fp16_enabled") and args[0].fp16_enabled): return old_func(*args, **kwargs) # get the arg spec of the decorated method @@ -105,11 +99,10 @@ def new_func(*args, **kwargs): new_args = [] # NOTE: default args are not taken into consideration if args: - arg_names = args_info.args[:len(args)] + arg_names = args_info.args[: len(args)] for i, arg_name in enumerate(arg_names): if arg_name in args_to_cast: - new_args.append( - cast_tensor_type(args[i], torch.float, torch.half)) + new_args.append(cast_tensor_type(args[i], torch.float, torch.half)) else: new_args.append(args[i]) # convert the kwargs that need to be processed @@ -117,13 +110,11 @@ def new_func(*args, **kwargs): if kwargs: for arg_name, arg_value in kwargs.items(): if arg_name in args_to_cast: - new_kwargs[arg_name] = cast_tensor_type( - arg_value, torch.float, torch.half) + new_kwargs[arg_name] = cast_tensor_type(arg_value, torch.float, torch.half) else: new_kwargs[arg_name] = arg_value # apply converted arguments to the decorated method - if (TORCH_VERSION != 'parrots' and - digit_version(TORCH_VERSION) >= digit_version('1.6.0')): + if TORCH_VERSION != "parrots" and digit_version(TORCH_VERSION) >= digit_version("1.6.0"): with autocast(enabled=True): output = old_func(*new_args, **new_kwargs) else: @@ -174,15 +165,13 @@ def force_fp32(apply_to=None, out_fp16=False): """ def force_fp32_wrapper(old_func): - @functools.wraps(old_func) def new_func(*args, **kwargs): # check if the module has set the attribute `fp16_enabled`, if not, # just fallback to the original method. if not isinstance(args[0], torch.nn.Module): - raise TypeError('@force_fp32 can only be used to decorate the ' - 'method of nn.Module') - if not (hasattr(args[0], 'fp16_enabled') and args[0].fp16_enabled): + raise TypeError("@force_fp32 can only be used to decorate the " "method of nn.Module") + if not (hasattr(args[0], "fp16_enabled") and args[0].fp16_enabled): return old_func(*args, **kwargs) # get the arg spec of the decorated method args_info = getfullargspec(old_func) @@ -191,11 +180,10 @@ def new_func(*args, **kwargs): # convert the args that need to be processed new_args = [] if args: - arg_names = args_info.args[:len(args)] + arg_names = args_info.args[: len(args)] for i, arg_name in enumerate(arg_names): if arg_name in args_to_cast: - new_args.append( - cast_tensor_type(args[i], torch.half, torch.float)) + new_args.append(cast_tensor_type(args[i], torch.half, torch.float)) else: new_args.append(args[i]) # convert the kwargs that need to be processed @@ -203,13 +191,11 @@ def new_func(*args, **kwargs): if kwargs: for arg_name, arg_value in kwargs.items(): if arg_name in args_to_cast: - new_kwargs[arg_name] = cast_tensor_type( - arg_value, torch.half, torch.float) + new_kwargs[arg_name] = cast_tensor_type(arg_value, torch.half, torch.float) else: new_kwargs[arg_name] = arg_value # apply converted arguments to the decorated method - if (TORCH_VERSION != 'parrots' and - digit_version(TORCH_VERSION) >= digit_version('1.6.0')): + if TORCH_VERSION != "parrots" and digit_version(TORCH_VERSION) >= digit_version("1.6.0"): with autocast(enabled=False): output = old_func(*new_args, **new_kwargs) else: @@ -227,7 +213,8 @@ def new_func(*args, **kwargs): def allreduce_grads(params, coalesce=True, bucket_size_mb=-1): warnings.warning( '"mmcv.runner.fp16_utils.allreduce_grads" is deprecated, and will be ' - 'removed in v2.8. Please switch to "mmcv.runner.allreduce_grads') + 'removed in v2.8. Please switch to "mmcv.runner.allreduce_grads' + ) _allreduce_grads(params, coalesce=coalesce, bucket_size_mb=bucket_size_mb) @@ -248,15 +235,14 @@ def wrap_fp16_model(model): Args: model (nn.Module): Model in FP32. """ - if (TORCH_VERSION == 'parrots' - or digit_version(TORCH_VERSION) < digit_version('1.6.0')): + if TORCH_VERSION == "parrots" or digit_version(TORCH_VERSION) < digit_version("1.6.0"): # convert model to fp16 model.half() # patch the normalization layers to make it work in fp32 mode patch_norm_fp32(model) # set `fp16_enabled` flag for m in model.modules(): - if hasattr(m, 'fp16_enabled'): + if hasattr(m, "fp16_enabled"): m.fp16_enabled = True @@ -272,9 +258,8 @@ def patch_norm_fp32(module): """ if isinstance(module, (nn.modules.batchnorm._BatchNorm, nn.GroupNorm)): module.float() - if isinstance(module, nn.GroupNorm) or torch.__version__ < '1.3': - module.forward = patch_forward_method(module.forward, torch.half, - torch.float) + if isinstance(module, nn.GroupNorm) or torch.__version__ < "1.3": + module.forward = patch_forward_method(module.forward, torch.half, torch.float) for child in module.children(): patch_norm_fp32(child) return module @@ -294,8 +279,7 @@ def patch_forward_method(func, src_type, dst_type, convert_output=True): """ def new_forward(*args, **kwargs): - output = func(*cast_tensor_type(args, src_type, dst_type), - **cast_tensor_type(kwargs, src_type, dst_type)) + output = func(*cast_tensor_type(args, src_type, dst_type), **cast_tensor_type(kwargs, src_type, dst_type)) if convert_output: output = cast_tensor_type(output, dst_type, src_type) return output @@ -332,15 +316,10 @@ class LossScaler: overflow to wait before increasing the loss scale. Default: 1000. """ - def __init__(self, - init_scale=2**32, - mode='dynamic', - scale_factor=2., - scale_window=1000): + def __init__(self, init_scale=2**32, mode="dynamic", scale_factor=2.0, scale_window=1000): self.cur_scale = init_scale self.cur_iter = 0 - assert mode in ('dynamic', - 'static'), 'mode can only be dynamic or static' + assert mode in ("dynamic", "static"), "mode can only be dynamic or static" self.mode = mode self.last_overflow_iter = -1 self.scale_factor = scale_factor @@ -348,7 +327,7 @@ def __init__(self, def has_overflow(self, params): """Check if params contain overflow.""" - if self.mode != 'dynamic': + if self.mode != "dynamic": return False for p in params: if p.grad is not None and LossScaler._has_inf_or_nan(p.grad.data): @@ -360,25 +339,23 @@ def _has_inf_or_nan(x): try: cpu_sum = float(x.float().sum()) except RuntimeError as instance: - if 'value cannot be converted' not in instance.args[0]: + if "value cannot be converted" not in instance.args[0]: raise return True else: - if cpu_sum == float('inf') or cpu_sum == -float('inf') \ - or cpu_sum != cpu_sum: + if cpu_sum == float("inf") or cpu_sum == -float("inf") or cpu_sum != cpu_sum: return True return False def update_scale(self, overflow): """update the current loss scale value when overflow happens.""" - if self.mode != 'dynamic': + if self.mode != "dynamic": return if overflow: self.cur_scale = max(self.cur_scale / self.scale_factor, 1) self.last_overflow_iter = self.cur_iter else: - if (self.cur_iter - self.last_overflow_iter) % \ - self.scale_window == 0: + if (self.cur_iter - self.last_overflow_iter) % self.scale_window == 0: self.cur_scale *= self.scale_factor self.cur_iter += 1 @@ -390,7 +367,8 @@ def state_dict(self): mode=self.mode, last_overflow_iter=self.last_overflow_iter, scale_factor=self.scale_factor, - scale_window=self.scale_window) + scale_window=self.scale_window, + ) def load_state_dict(self, state_dict): """Loads the loss_scaler state dict. @@ -398,12 +376,12 @@ def load_state_dict(self, state_dict): Args: state_dict (dict): scaler state. """ - self.cur_scale = state_dict['cur_scale'] - self.cur_iter = state_dict['cur_iter'] - self.mode = state_dict['mode'] - self.last_overflow_iter = state_dict['last_overflow_iter'] - self.scale_factor = state_dict['scale_factor'] - self.scale_window = state_dict['scale_window'] + self.cur_scale = state_dict["cur_scale"] + self.cur_iter = state_dict["cur_iter"] + self.mode = state_dict["mode"] + self.last_overflow_iter = state_dict["last_overflow_iter"] + self.scale_factor = state_dict["scale_factor"] + self.scale_window = state_dict["scale_window"] @property def loss_scale(self): diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/checkpoint.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/checkpoint.py index 816ea17c..7bd373a9 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/checkpoint.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/checkpoint.py @@ -2,7 +2,7 @@ import os.path as osp import warnings -from comfy_controlnet_preprocessors.uniformer.mmcv.fileio import FileClient +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.fileio import FileClient from ..dist_utils import allreduce_params, master_only from .hook import HOOKS, Hook @@ -48,16 +48,18 @@ class CheckpointHook(Hook): path will be "/path/of/A/B". """ - def __init__(self, - interval=-1, - by_epoch=True, - save_optimizer=True, - out_dir=None, - max_keep_ckpts=-1, - save_last=True, - sync_buffer=False, - file_client_args=None, - **kwargs): + def __init__( + self, + interval=-1, + by_epoch=True, + save_optimizer=True, + out_dir=None, + max_keep_ckpts=-1, + save_last=True, + sync_buffer=False, + file_client_args=None, + **kwargs, + ): self.interval = interval self.by_epoch = by_epoch self.save_optimizer = save_optimizer @@ -72,8 +74,7 @@ def before_run(self, runner): if not self.out_dir: self.out_dir = runner.work_dir - self.file_client = FileClient.infer_client(self.file_client_args, - self.out_dir) + self.file_client = FileClient.infer_client(self.file_client_args, self.out_dir) # if `self.out_dir` is not equal to `runner.work_dir`, it means that # `self.out_dir` is set so the final `self.out_dir` is the @@ -83,21 +84,22 @@ def before_run(self, runner): basename = osp.basename(runner.work_dir.rstrip(osp.sep)) self.out_dir = self.file_client.join_path(self.out_dir, basename) - runner.logger.info((f'Checkpoints will be saved to {self.out_dir} by ' - f'{self.file_client.name}.')) + runner.logger.info((f"Checkpoints will be saved to {self.out_dir} by " f"{self.file_client.name}.")) # disable the create_symlink option because some file backends do not # allow to create a symlink - if 'create_symlink' in self.args: - if self.args[ - 'create_symlink'] and not self.file_client.allow_symlink: - self.args['create_symlink'] = False + if "create_symlink" in self.args: + if self.args["create_symlink"] and not self.file_client.allow_symlink: + self.args["create_symlink"] = False warnings.warn( - ('create_symlink is set as True by the user but is changed' - 'to be False because creating symbolic link is not ' - f'allowed in {self.file_client.name}')) + ( + "create_symlink is set as True by the user but is changed" + "to be False because creating symbolic link is not " + f"allowed in {self.file_client.name}" + ) + ) else: - self.args['create_symlink'] = self.file_client.allow_symlink + self.args["create_symlink"] = self.file_client.allow_symlink def after_train_epoch(self, runner): if not self.by_epoch: @@ -106,11 +108,8 @@ def after_train_epoch(self, runner): # save checkpoint for following cases: # 1. every ``self.interval`` epochs # 2. reach the last epoch of training - if self.every_n_epochs( - runner, self.interval) or (self.save_last - and self.is_last_epoch(runner)): - runner.logger.info( - f'Saving checkpoint at {runner.epoch + 1} epochs') + if self.every_n_epochs(runner, self.interval) or (self.save_last and self.is_last_epoch(runner)): + runner.logger.info(f"Saving checkpoint at {runner.epoch + 1} epochs") if self.sync_buffer: allreduce_params(runner.model.buffers()) self._save_checkpoint(runner) @@ -118,33 +117,26 @@ def after_train_epoch(self, runner): @master_only def _save_checkpoint(self, runner): """Save the current checkpoint and delete unwanted checkpoint.""" - runner.save_checkpoint( - self.out_dir, save_optimizer=self.save_optimizer, **self.args) + runner.save_checkpoint(self.out_dir, save_optimizer=self.save_optimizer, **self.args) if runner.meta is not None: if self.by_epoch: - cur_ckpt_filename = self.args.get( - 'filename_tmpl', 'epoch_{}.pth').format(runner.epoch + 1) + cur_ckpt_filename = self.args.get("filename_tmpl", "epoch_{}.pth").format(runner.epoch + 1) else: - cur_ckpt_filename = self.args.get( - 'filename_tmpl', 'iter_{}.pth').format(runner.iter + 1) - runner.meta.setdefault('hook_msgs', dict()) - runner.meta['hook_msgs']['last_ckpt'] = self.file_client.join_path( - self.out_dir, cur_ckpt_filename) + cur_ckpt_filename = self.args.get("filename_tmpl", "iter_{}.pth").format(runner.iter + 1) + runner.meta.setdefault("hook_msgs", dict()) + runner.meta["hook_msgs"]["last_ckpt"] = self.file_client.join_path(self.out_dir, cur_ckpt_filename) # remove other checkpoints if self.max_keep_ckpts > 0: if self.by_epoch: - name = 'epoch_{}.pth' + name = "epoch_{}.pth" current_ckpt = runner.epoch + 1 else: - name = 'iter_{}.pth' + name = "iter_{}.pth" current_ckpt = runner.iter + 1 - redundant_ckpts = range( - current_ckpt - self.max_keep_ckpts * self.interval, 0, - -self.interval) - filename_tmpl = self.args.get('filename_tmpl', name) + redundant_ckpts = range(current_ckpt - self.max_keep_ckpts * self.interval, 0, -self.interval) + filename_tmpl = self.args.get("filename_tmpl", name) for _step in redundant_ckpts: - ckpt_path = self.file_client.join_path( - self.out_dir, filename_tmpl.format(_step)) + ckpt_path = self.file_client.join_path(self.out_dir, filename_tmpl.format(_step)) if self.file_client.isfile(ckpt_path): self.file_client.remove(ckpt_path) else: @@ -157,11 +149,8 @@ def after_train_iter(self, runner): # save checkpoint for following cases: # 1. every ``self.interval`` iterations # 2. reach the last iteration of training - if self.every_n_iters( - runner, self.interval) or (self.save_last - and self.is_last_iter(runner)): - runner.logger.info( - f'Saving checkpoint at {runner.iter + 1} iterations') + if self.every_n_iters(runner, self.interval) or (self.save_last and self.is_last_iter(runner)): + runner.logger.info(f"Saving checkpoint at {runner.iter + 1} iterations") if self.sync_buffer: allreduce_params(runner.model.buffers()) self._save_checkpoint(runner) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/evaluation.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/evaluation.py index ca72c223..8cfb5ec7 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/evaluation.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/evaluation.py @@ -7,8 +7,8 @@ from torch.nn.modules.batchnorm import _BatchNorm from torch.utils.data import DataLoader -from comfy_controlnet_preprocessors.uniformer.mmcv.fileio import FileClient -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import is_seq_of +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.fileio import FileClient +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import is_seq_of from .hook import Hook from .logger import LoggerHook @@ -74,55 +74,52 @@ class EvalHook(Hook): # tasks, downstream repos may need to overwrite the following inner # variable accordingly. - rule_map = {'greater': lambda x, y: x > y, 'less': lambda x, y: x < y} - init_value_map = {'greater': -inf, 'less': inf} - _default_greater_keys = [ - 'acc', 'top', 'AR@', 'auc', 'precision', 'mAP', 'mDice', 'mIoU', - 'mAcc', 'aAcc' - ] - _default_less_keys = ['loss'] - - def __init__(self, - dataloader, - start=None, - interval=1, - by_epoch=True, - save_best=None, - rule=None, - test_fn=None, - greater_keys=None, - less_keys=None, - out_dir=None, - file_client_args=None, - **eval_kwargs): + rule_map = {"greater": lambda x, y: x > y, "less": lambda x, y: x < y} + init_value_map = {"greater": -inf, "less": inf} + _default_greater_keys = ["acc", "top", "AR@", "auc", "precision", "mAP", "mDice", "mIoU", "mAcc", "aAcc"] + _default_less_keys = ["loss"] + + def __init__( + self, + dataloader, + start=None, + interval=1, + by_epoch=True, + save_best=None, + rule=None, + test_fn=None, + greater_keys=None, + less_keys=None, + out_dir=None, + file_client_args=None, + **eval_kwargs, + ): if not isinstance(dataloader, DataLoader): - raise TypeError(f'dataloader must be a pytorch DataLoader, ' - f'but got {type(dataloader)}') + raise TypeError(f"dataloader must be a pytorch DataLoader, " f"but got {type(dataloader)}") if interval <= 0: - raise ValueError(f'interval must be a positive number, ' - f'but got {interval}') + raise ValueError(f"interval must be a positive number, " f"but got {interval}") - assert isinstance(by_epoch, bool), '``by_epoch`` should be a boolean' + assert isinstance(by_epoch, bool), "``by_epoch`` should be a boolean" if start is not None and start < 0: - raise ValueError(f'The evaluation start epoch {start} is smaller ' - f'than 0') + raise ValueError(f"The evaluation start epoch {start} is smaller " f"than 0") self.dataloader = dataloader self.interval = interval self.start = start self.by_epoch = by_epoch - assert isinstance(save_best, str) or save_best is None, \ - '""save_best"" should be a str or None ' \ - f'rather than {type(save_best)}' + assert isinstance(save_best, str) or save_best is None, ( + '""save_best"" should be a str or None ' f"rather than {type(save_best)}" + ) self.save_best = save_best self.eval_kwargs = eval_kwargs self.initial_flag = True if test_fn is None: - from comfy_controlnet_preprocessors.uniformer.mmcv.engine import single_gpu_test + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.engine import single_gpu_test + self.test_fn = single_gpu_test else: self.test_fn = test_fn @@ -131,7 +128,7 @@ def __init__(self, self.greater_keys = self._default_greater_keys else: if not isinstance(greater_keys, (list, tuple)): - greater_keys = (greater_keys, ) + greater_keys = (greater_keys,) assert is_seq_of(greater_keys, str) self.greater_keys = greater_keys @@ -139,7 +136,7 @@ def __init__(self, self.less_keys = self._default_less_keys else: if not isinstance(less_keys, (list, tuple)): - less_keys = (less_keys, ) + less_keys = (less_keys,) assert is_seq_of(less_keys, str) self.less_keys = less_keys @@ -171,11 +168,10 @@ def _init_rule(self, rule, key_indicator): comparison rule. """ if rule not in self.rule_map and rule is not None: - raise KeyError(f'rule must be greater, less or None, ' - f'but got {rule}.') + raise KeyError(f"rule must be greater, less or None, " f"but got {rule}.") if rule is None: - if key_indicator != 'auto': + if key_indicator != "auto": # `_lc` here means we use the lower case of keys for # case-insensitive matching key_indicator_lc = key_indicator.lower() @@ -183,17 +179,19 @@ def _init_rule(self, rule, key_indicator): less_keys = [key.lower() for key in self.less_keys] if key_indicator_lc in greater_keys: - rule = 'greater' + rule = "greater" elif key_indicator_lc in less_keys: - rule = 'less' + rule = "less" elif any(key in key_indicator_lc for key in greater_keys): - rule = 'greater' + rule = "greater" elif any(key in key_indicator_lc for key in less_keys): - rule = 'less' + rule = "less" else: - raise ValueError(f'Cannot infer the rule for key ' - f'{key_indicator}, thus a specific rule ' - f'must be specified.') + raise ValueError( + f"Cannot infer the rule for key " + f"{key_indicator}, thus a specific rule " + f"must be specified." + ) self.rule = rule self.key_indicator = key_indicator if self.rule is not None: @@ -203,8 +201,7 @@ def before_run(self, runner): if not self.out_dir: self.out_dir = runner.work_dir - self.file_client = FileClient.infer_client(self.file_client_args, - self.out_dir) + self.file_client = FileClient.infer_client(self.file_client_args, self.out_dir) # if `self.out_dir` is not equal to `runner.work_dir`, it means that # `self.out_dir` is set so the final `self.out_dir` is the @@ -213,17 +210,14 @@ def before_run(self, runner): if self.out_dir != runner.work_dir: basename = osp.basename(runner.work_dir.rstrip(osp.sep)) self.out_dir = self.file_client.join_path(self.out_dir, basename) - runner.logger.info( - (f'The best checkpoint will be saved to {self.out_dir} by ' - f'{self.file_client.name}')) + runner.logger.info((f"The best checkpoint will be saved to {self.out_dir} by " f"{self.file_client.name}")) if self.save_best is not None: if runner.meta is None: - warnings.warn('runner.meta is None. Creating an empty one.') + warnings.warn("runner.meta is None. Creating an empty one.") runner.meta = dict() - runner.meta.setdefault('hook_msgs', dict()) - self.best_ckpt_path = runner.meta['hook_msgs'].get( - 'best_ckpt', None) + runner.meta.setdefault("hook_msgs", dict()) + self.best_ckpt_path = runner.meta["hook_msgs"].get("best_ckpt", None) def before_train_iter(self, runner): """Evaluate the model only at the start of training by iteration.""" @@ -269,7 +263,7 @@ def after_train_epoch(self, runner): def _do_evaluate(self, runner): """perform evaluation and save ckpt.""" results = self.test_fn(runner.model, self.dataloader) - runner.log_buffer.output['eval_iter_num'] = len(self.dataloader) + runner.log_buffer.output["eval_iter_num"] = len(self.dataloader) key_score = self.evaluate(runner, results) # the key_score may be `None` so it needs to skip the action to save # the best checkpoint @@ -319,37 +313,28 @@ def _save_ckpt(self, runner, key_score): best checkpoint into ``work_dir``. """ if self.by_epoch: - current = f'epoch_{runner.epoch + 1}' - cur_type, cur_time = 'epoch', runner.epoch + 1 + current = f"epoch_{runner.epoch + 1}" + cur_type, cur_time = "epoch", runner.epoch + 1 else: - current = f'iter_{runner.iter + 1}' - cur_type, cur_time = 'iter', runner.iter + 1 + current = f"iter_{runner.iter + 1}" + cur_type, cur_time = "iter", runner.iter + 1 - best_score = runner.meta['hook_msgs'].get( - 'best_score', self.init_value_map[self.rule]) + best_score = runner.meta["hook_msgs"].get("best_score", self.init_value_map[self.rule]) if self.compare_func(key_score, best_score): best_score = key_score - runner.meta['hook_msgs']['best_score'] = best_score + runner.meta["hook_msgs"]["best_score"] = best_score - if self.best_ckpt_path and self.file_client.isfile( - self.best_ckpt_path): + if self.best_ckpt_path and self.file_client.isfile(self.best_ckpt_path): self.file_client.remove(self.best_ckpt_path) - runner.logger.info( - (f'The previous best checkpoint {self.best_ckpt_path} was ' - 'removed')) - - best_ckpt_name = f'best_{self.key_indicator}_{current}.pth' - self.best_ckpt_path = self.file_client.join_path( - self.out_dir, best_ckpt_name) - runner.meta['hook_msgs']['best_ckpt'] = self.best_ckpt_path - - runner.save_checkpoint( - self.out_dir, best_ckpt_name, create_symlink=False) - runner.logger.info( - f'Now best checkpoint is saved as {best_ckpt_name}.') - runner.logger.info( - f'Best {self.key_indicator} is {best_score:0.4f} ' - f'at {cur_time} {cur_type}.') + runner.logger.info((f"The previous best checkpoint {self.best_ckpt_path} was " "removed")) + + best_ckpt_name = f"best_{self.key_indicator}_{current}.pth" + self.best_ckpt_path = self.file_client.join_path(self.out_dir, best_ckpt_name) + runner.meta["hook_msgs"]["best_ckpt"] = self.best_ckpt_path + + runner.save_checkpoint(self.out_dir, best_ckpt_name, create_symlink=False) + runner.logger.info(f"Now best checkpoint is saved as {best_ckpt_name}.") + runner.logger.info(f"Best {self.key_indicator} is {best_score:0.4f} " f"at {cur_time} {cur_type}.") def evaluate(self, runner, results): """Evaluate the results. @@ -358,8 +343,7 @@ def evaluate(self, runner, results): runner (:obj:`mmcv.Runner`): The underlined training runner. results (list): Output results. """ - eval_res = self.dataloader.dataset.evaluate( - results, logger=runner.logger, **self.eval_kwargs) + eval_res = self.dataloader.dataset.evaluate(results, logger=runner.logger, **self.eval_kwargs) for name, val in eval_res.items(): runner.log_buffer.output[name] = val @@ -372,11 +356,12 @@ def evaluate(self, runner, results): # https://github.com/open-mmlab/mmdetection/issues/6265. if not eval_res: warnings.warn( - 'Since `eval_res` is an empty dict, the behavior to save ' - 'the best checkpoint will be skipped in this evaluation.') + "Since `eval_res` is an empty dict, the behavior to save " + "the best checkpoint will be skipped in this evaluation." + ) return None - if self.key_indicator == 'auto': + if self.key_indicator == "auto": # infer from eval_results self._init_rule(self.rule, list(eval_res.keys())[0]) return eval_res[self.key_indicator] @@ -436,25 +421,27 @@ class DistEvalHook(EvalHook): the dataset. """ - def __init__(self, - dataloader, - start=None, - interval=1, - by_epoch=True, - save_best=None, - rule=None, - test_fn=None, - greater_keys=None, - less_keys=None, - broadcast_bn_buffer=True, - tmpdir=None, - gpu_collect=False, - out_dir=None, - file_client_args=None, - **eval_kwargs): - + def __init__( + self, + dataloader, + start=None, + interval=1, + by_epoch=True, + save_best=None, + rule=None, + test_fn=None, + greater_keys=None, + less_keys=None, + broadcast_bn_buffer=True, + tmpdir=None, + gpu_collect=False, + out_dir=None, + file_client_args=None, + **eval_kwargs, + ): if test_fn is None: - from comfy_controlnet_preprocessors.uniformer.mmcv.engine import multi_gpu_test + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.engine import multi_gpu_test + test_fn = multi_gpu_test super().__init__( @@ -469,7 +456,8 @@ def __init__(self, less_keys=less_keys, out_dir=out_dir, file_client_args=file_client_args, - **eval_kwargs) + **eval_kwargs, + ) self.broadcast_bn_buffer = broadcast_bn_buffer self.tmpdir = tmpdir @@ -485,23 +473,18 @@ def _do_evaluate(self, runner): if self.broadcast_bn_buffer: model = runner.model for name, module in model.named_modules(): - if isinstance(module, - _BatchNorm) and module.track_running_stats: + if isinstance(module, _BatchNorm) and module.track_running_stats: dist.broadcast(module.running_var, 0) dist.broadcast(module.running_mean, 0) tmpdir = self.tmpdir if tmpdir is None: - tmpdir = osp.join(runner.work_dir, '.eval_hook') + tmpdir = osp.join(runner.work_dir, ".eval_hook") - results = self.test_fn( - runner.model, - self.dataloader, - tmpdir=tmpdir, - gpu_collect=self.gpu_collect) + results = self.test_fn(runner.model, self.dataloader, tmpdir=tmpdir, gpu_collect=self.gpu_collect) if runner.rank == 0: - print('\n') - runner.log_buffer.output['eval_iter_num'] = len(self.dataloader) + print("\n") + runner.log_buffer.output["eval_iter_num"] = len(self.dataloader) key_score = self.evaluate(runner, results) # the key_score may be `None` so it needs to skip the action to # save the best checkpoint diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/hook.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/hook.py index 4b618a02..72c62778 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/hook.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/hook.py @@ -1,6 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved.\ import sys -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry, is_method_overridden +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry, is_method_overridden HOOKS = Registry("hook", sys.modules[__name__]) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/pavi.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/pavi.py index 9580e7b8..358c60ea 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/pavi.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/pavi.py @@ -6,7 +6,7 @@ import torch import yaml -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv from ....parallel.utils import is_module_wrapper from ...dist_utils import master_only from ..hook import HOOKS @@ -15,18 +15,18 @@ @HOOKS.register_module() class PaviLoggerHook(LoggerHook): - - def __init__(self, - init_kwargs=None, - add_graph=False, - add_last_ckpt=False, - interval=10, - ignore_last=True, - reset_flag=False, - by_epoch=True, - img_key='img_info'): - super(PaviLoggerHook, self).__init__(interval, ignore_last, reset_flag, - by_epoch) + def __init__( + self, + init_kwargs=None, + add_graph=False, + add_last_ckpt=False, + interval=10, + ignore_last=True, + reset_flag=False, + by_epoch=True, + img_key="img_info", + ): + super(PaviLoggerHook, self).__init__(interval, ignore_last, reset_flag, by_epoch) self.init_kwargs = init_kwargs self.add_graph = add_graph self.add_last_ckpt = add_last_ckpt @@ -40,21 +40,20 @@ def before_run(self, runner): except ImportError: raise ImportError('Please run "pip install pavi" to install pavi.') - self.run_name = runner.work_dir.split('/')[-1] + self.run_name = runner.work_dir.split("/")[-1] if not self.init_kwargs: self.init_kwargs = dict() - self.init_kwargs['name'] = self.run_name - self.init_kwargs['model'] = runner._model_name + self.init_kwargs["name"] = self.run_name + self.init_kwargs["model"] = runner._model_name if runner.meta is not None: - if 'config_dict' in runner.meta: - config_dict = runner.meta['config_dict'] - assert isinstance( - config_dict, - dict), ('meta["config_dict"] has to be of a dict, ' - f'but got {type(config_dict)}') - elif 'config_file' in runner.meta: - config_file = runner.meta['config_file'] + if "config_dict" in runner.meta: + config_dict = runner.meta["config_dict"] + assert isinstance(config_dict, dict), ( + 'meta["config_dict"] has to be of a dict, ' f"but got {type(config_dict)}" + ) + elif "config_file" in runner.meta: + config_file = runner.meta["config_file"] config_dict = dict(mmcv.Config.fromfile(config_file)) else: config_dict = None @@ -62,18 +61,17 @@ def before_run(self, runner): # 'max_.*iter' is parsed in pavi sdk as the maximum iterations # to properly set up the progress bar. config_dict = config_dict.copy() - config_dict.setdefault('max_iter', runner.max_iters) + config_dict.setdefault("max_iter", runner.max_iters) # non-serializable values are first converted in # mmcv.dump to json - config_dict = json.loads( - mmcv.dump(config_dict, file_format='json')) + config_dict = json.loads(mmcv.dump(config_dict, file_format="json")) session_text = yaml.dump(config_dict) - self.init_kwargs['session_text'] = session_text + self.init_kwargs["session_text"] = session_text self.writer = SummaryWriter(**self.init_kwargs) def get_step(self, runner): """Get the total training step/epoch.""" - if self.get_mode(runner) == 'val' and self.by_epoch: + if self.get_mode(runner) == "val" and self.by_epoch: return self.get_epoch(runner) else: return self.get_iter(runner) @@ -82,13 +80,12 @@ def get_step(self, runner): def log(self, runner): tags = self.get_loggable_tags(runner, add_mode=False) if tags: - self.writer.add_scalars( - self.get_mode(runner), tags, self.get_step(runner)) + self.writer.add_scalars(self.get_mode(runner), tags, self.get_step(runner)) @master_only def after_run(self, runner): if self.add_last_ckpt: - ckpt_path = osp.join(runner.work_dir, 'latest.pth') + ckpt_path = osp.join(runner.work_dir, "latest.pth") if osp.islink(ckpt_path): ckpt_path = osp.join(runner.work_dir, os.readlink(ckpt_path)) @@ -96,9 +93,8 @@ def after_run(self, runner): # runner.epoch += 1 has been done before `after_run`. iteration = runner.epoch if self.by_epoch else runner.iter return self.writer.add_snapshot_file( - tag=self.run_name, - snapshot_file_path=ckpt_path, - iteration=iteration) + tag=self.run_name, snapshot_file_path=ckpt_path, iteration=iteration + ) # flush the buffer and send a task ending signal to Pavi self.writer.close() diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/tensorboard.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/tensorboard.py index 69887228..9b8f2ab7 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/tensorboard.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/tensorboard.py @@ -1,7 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. import os.path as osp -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, digit_version +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, digit_version from ...dist_utils import master_only from ..hook import HOOKS from .base import LoggerHook @@ -9,38 +9,30 @@ @HOOKS.register_module() class TensorboardLoggerHook(LoggerHook): - - def __init__(self, - log_dir=None, - interval=10, - ignore_last=True, - reset_flag=False, - by_epoch=True): - super(TensorboardLoggerHook, self).__init__(interval, ignore_last, - reset_flag, by_epoch) + def __init__(self, log_dir=None, interval=10, ignore_last=True, reset_flag=False, by_epoch=True): + super(TensorboardLoggerHook, self).__init__(interval, ignore_last, reset_flag, by_epoch) self.log_dir = log_dir @master_only def before_run(self, runner): super(TensorboardLoggerHook, self).before_run(runner) - if (TORCH_VERSION == 'parrots' - or digit_version(TORCH_VERSION) < digit_version('1.1')): + if TORCH_VERSION == "parrots" or digit_version(TORCH_VERSION) < digit_version("1.1"): try: from tensorboardX import SummaryWriter except ImportError: - raise ImportError('Please install tensorboardX to use ' - 'TensorboardLoggerHook.') + raise ImportError("Please install tensorboardX to use " "TensorboardLoggerHook.") else: try: from torch.utils.tensorboard import SummaryWriter except ImportError: raise ImportError( 'Please run "pip install future tensorboard" to install ' - 'the dependencies to use torch.utils.tensorboard ' - '(applicable to PyTorch 1.1 or higher)') + "the dependencies to use torch.utils.tensorboard " + "(applicable to PyTorch 1.1 or higher)" + ) if self.log_dir is None: - self.log_dir = osp.join(runner.work_dir, 'tf_logs') + self.log_dir = osp.join(runner.work_dir, "tf_logs") self.writer = SummaryWriter(self.log_dir) @master_only diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/text.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/text.py index d934b7b7..d6e8dd7e 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/text.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/logger/text.py @@ -7,9 +7,9 @@ import torch import torch.distributed as dist -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv -from comfy_controlnet_preprocessors.uniformer.mmcv.fileio.file_client import FileClient -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import is_tuple_of, scandir +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.fileio.file_client import FileClient +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import is_tuple_of, scandir from ..hook import HOOKS from .base import LoggerHook @@ -52,66 +52,61 @@ class TextLoggerHook(LoggerHook): `New in version 1.3.16.` """ - def __init__(self, - by_epoch=True, - interval=10, - ignore_last=True, - reset_flag=False, - interval_exp_name=1000, - out_dir=None, - out_suffix=('.log.json', '.log', '.py'), - keep_local=True, - file_client_args=None): - super(TextLoggerHook, self).__init__(interval, ignore_last, reset_flag, - by_epoch) + def __init__( + self, + by_epoch=True, + interval=10, + ignore_last=True, + reset_flag=False, + interval_exp_name=1000, + out_dir=None, + out_suffix=(".log.json", ".log", ".py"), + keep_local=True, + file_client_args=None, + ): + super(TextLoggerHook, self).__init__(interval, ignore_last, reset_flag, by_epoch) self.by_epoch = by_epoch self.time_sec_tot = 0 self.interval_exp_name = interval_exp_name if out_dir is None and file_client_args is not None: - raise ValueError( - 'file_client_args should be "None" when `out_dir` is not' - 'specified.') + raise ValueError('file_client_args should be "None" when `out_dir` is not' "specified.") self.out_dir = out_dir - if not (out_dir is None or isinstance(out_dir, str) - or is_tuple_of(out_dir, str)): - raise TypeError('out_dir should be "None" or string or tuple of ' - 'string, but got {out_dir}') + if not (out_dir is None or isinstance(out_dir, str) or is_tuple_of(out_dir, str)): + raise TypeError('out_dir should be "None" or string or tuple of ' "string, but got {out_dir}") self.out_suffix = out_suffix self.keep_local = keep_local self.file_client_args = file_client_args if self.out_dir is not None: - self.file_client = FileClient.infer_client(file_client_args, - self.out_dir) + self.file_client = FileClient.infer_client(file_client_args, self.out_dir) def before_run(self, runner): super(TextLoggerHook, self).before_run(runner) if self.out_dir is not None: - self.file_client = FileClient.infer_client(self.file_client_args, - self.out_dir) + self.file_client = FileClient.infer_client(self.file_client_args, self.out_dir) # The final `self.out_dir` is the concatenation of `self.out_dir` # and the last level directory of `runner.work_dir` basename = osp.basename(runner.work_dir.rstrip(osp.sep)) self.out_dir = self.file_client.join_path(self.out_dir, basename) runner.logger.info( - (f'Text logs will be saved to {self.out_dir} by ' - f'{self.file_client.name} after the training process.')) + ( + f"Text logs will be saved to {self.out_dir} by " + f"{self.file_client.name} after the training process." + ) + ) self.start_iter = runner.iter - self.json_log_path = osp.join(runner.work_dir, - f'{runner.timestamp}.log.json') + self.json_log_path = osp.join(runner.work_dir, f"{runner.timestamp}.log.json") if runner.meta is not None: self._dump_log(runner.meta, runner) def _get_max_memory(self, runner): - device = getattr(runner.model, 'output_device', None) + device = getattr(runner.model, "output_device", None) mem = torch.cuda.max_memory_allocated(device=device) - mem_mb = torch.tensor([mem / (1024 * 1024)], - dtype=torch.int, - device=device) + mem_mb = torch.tensor([mem / (1024 * 1024)], dtype=torch.int, device=device) if runner.world_size > 1: dist.reduce(mem_mb, 0, op=dist.ReduceOp.MAX) return mem_mb.item() @@ -119,39 +114,35 @@ def _get_max_memory(self, runner): def _log_info(self, log_dict, runner): # print exp name for users to distinguish experiments # at every ``interval_exp_name`` iterations and the end of each epoch - if runner.meta is not None and 'exp_name' in runner.meta: - if (self.every_n_iters(runner, self.interval_exp_name)) or ( - self.by_epoch and self.end_of_epoch(runner)): + if runner.meta is not None and "exp_name" in runner.meta: + if (self.every_n_iters(runner, self.interval_exp_name)) or (self.by_epoch and self.end_of_epoch(runner)): exp_info = f'Exp name: {runner.meta["exp_name"]}' runner.logger.info(exp_info) - if log_dict['mode'] == 'train': - if isinstance(log_dict['lr'], dict): + if log_dict["mode"] == "train": + if isinstance(log_dict["lr"], dict): lr_str = [] - for k, val in log_dict['lr'].items(): - lr_str.append(f'lr_{k}: {val:.3e}') - lr_str = ' '.join(lr_str) + for k, val in log_dict["lr"].items(): + lr_str.append(f"lr_{k}: {val:.3e}") + lr_str = " ".join(lr_str) else: lr_str = f'lr: {log_dict["lr"]:.3e}' # by epoch: Epoch [4][100/1000] # by iter: Iter [100/100000] if self.by_epoch: - log_str = f'Epoch [{log_dict["epoch"]}]' \ - f'[{log_dict["iter"]}/{len(runner.data_loader)}]\t' + log_str = f'Epoch [{log_dict["epoch"]}]' f'[{log_dict["iter"]}/{len(runner.data_loader)}]\t' else: log_str = f'Iter [{log_dict["iter"]}/{runner.max_iters}]\t' - log_str += f'{lr_str}, ' + log_str += f"{lr_str}, " - if 'time' in log_dict.keys(): - self.time_sec_tot += (log_dict['time'] * self.interval) - time_sec_avg = self.time_sec_tot / ( - runner.iter - self.start_iter + 1) + if "time" in log_dict.keys(): + self.time_sec_tot += log_dict["time"] * self.interval + time_sec_avg = self.time_sec_tot / (runner.iter - self.start_iter + 1) eta_sec = time_sec_avg * (runner.max_iters - runner.iter - 1) eta_str = str(datetime.timedelta(seconds=int(eta_sec))) - log_str += f'eta: {eta_str}, ' - log_str += f'time: {log_dict["time"]:.3f}, ' \ - f'data_time: {log_dict["data_time"]:.3f}, ' + log_str += f"eta: {eta_str}, " + log_str += f'time: {log_dict["time"]:.3f}, ' f'data_time: {log_dict["data_time"]:.3f}, ' # statistic memory if torch.cuda.is_available(): log_str += f'memory: {log_dict["memory"]}, ' @@ -161,8 +152,7 @@ def _log_info(self, log_dict, runner): # by epoch: Epoch[val] [4][1000] # by iter: Iter[val] [1000] if self.by_epoch: - log_str = f'Epoch({log_dict["mode"]}) ' \ - f'[{log_dict["epoch"]}][{log_dict["iter"]}]\t' + log_str = f'Epoch({log_dict["mode"]}) ' f'[{log_dict["epoch"]}][{log_dict["iter"]}]\t' else: log_str = f'Iter({log_dict["mode"]}) [{log_dict["iter"]}]\t' @@ -170,15 +160,12 @@ def _log_info(self, log_dict, runner): for name, val in log_dict.items(): # TODO: resolve this hack # these items have been in log_str - if name in [ - 'mode', 'Epoch', 'iter', 'lr', 'time', 'data_time', - 'memory', 'epoch' - ]: + if name in ["mode", "Epoch", "iter", "lr", "time", "data_time", "memory", "epoch"]: continue if isinstance(val, float): - val = f'{val:.4f}' - log_items.append(f'{name}: {val}') - log_str += ', '.join(log_items) + val = f"{val:.4f}" + log_items.append(f"{name}: {val}") + log_str += ", ".join(log_items) runner.logger.info(log_str) @@ -189,9 +176,9 @@ def _dump_log(self, log_dict, runner): json_log[k] = self._round_float(v) # only append log at last line if runner.rank == 0: - with open(self.json_log_path, 'a+') as f: - mmcv.dump(json_log, f, file_format='json') - f.write('\n') + with open(self.json_log_path, "a+") as f: + mmcv.dump(json_log, f, file_format="json") + f.write("\n") def _round_float(self, items): if isinstance(items, list): @@ -202,32 +189,29 @@ def _round_float(self, items): return items def log(self, runner): - if 'eval_iter_num' in runner.log_buffer.output: + if "eval_iter_num" in runner.log_buffer.output: # this doesn't modify runner.iter and is regardless of by_epoch - cur_iter = runner.log_buffer.output.pop('eval_iter_num') + cur_iter = runner.log_buffer.output.pop("eval_iter_num") else: cur_iter = self.get_iter(runner, inner_iter=True) - log_dict = OrderedDict( - mode=self.get_mode(runner), - epoch=self.get_epoch(runner), - iter=cur_iter) + log_dict = OrderedDict(mode=self.get_mode(runner), epoch=self.get_epoch(runner), iter=cur_iter) # only record lr of the first param group cur_lr = runner.current_lr() if isinstance(cur_lr, list): - log_dict['lr'] = cur_lr[0] + log_dict["lr"] = cur_lr[0] else: assert isinstance(cur_lr, dict) - log_dict['lr'] = {} + log_dict["lr"] = {} for k, lr_ in cur_lr.items(): assert isinstance(lr_, list) - log_dict['lr'].update({k: lr_[0]}) + log_dict["lr"].update({k: lr_[0]}) - if 'time' in runner.log_buffer.output: + if "time" in runner.log_buffer.output: # statistic memory if torch.cuda.is_available(): - log_dict['memory'] = self._get_max_memory(runner) + log_dict["memory"] = self._get_max_memory(runner) log_dict = dict(log_dict, **runner.log_buffer.output) @@ -240,17 +224,12 @@ def after_run(self, runner): if self.out_dir is not None: for filename in scandir(runner.work_dir, self.out_suffix, True): local_filepath = osp.join(runner.work_dir, filename) - out_filepath = self.file_client.join_path( - self.out_dir, filename) - with open(local_filepath, 'r') as f: + out_filepath = self.file_client.join_path(self.out_dir, filename) + with open(local_filepath, "r") as f: self.file_client.put_text(f.read(), out_filepath) - runner.logger.info( - (f'The file {local_filepath} has been uploaded to ' - f'{out_filepath}.')) + runner.logger.info((f"The file {local_filepath} has been uploaded to " f"{out_filepath}.")) if not self.keep_local: os.remove(local_filepath) - runner.logger.info( - (f'{local_filepath} was removed due to the ' - '`self.keep_local=False`')) + runner.logger.info((f"{local_filepath} was removed due to the " "`self.keep_local=False`")) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/lr_updater.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/lr_updater.py index a81993f5..5ee92f64 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/lr_updater.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/lr_updater.py @@ -2,7 +2,7 @@ import numbers from math import cos, pi -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv from .hook import HOOKS, Hook @@ -22,23 +22,16 @@ class LrUpdaterHook(Hook): number of iteration that warmup lasts """ - def __init__(self, - by_epoch=True, - warmup=None, - warmup_iters=0, - warmup_ratio=0.1, - warmup_by_epoch=False): + def __init__(self, by_epoch=True, warmup=None, warmup_iters=0, warmup_ratio=0.1, warmup_by_epoch=False): # validate the "warmup" argument if warmup is not None: - if warmup not in ['constant', 'linear', 'exp']: + if warmup not in ["constant", "linear", "exp"]: raise ValueError( - f'"{warmup}" is not a supported type for warming up, valid' - ' types are "constant" and "linear"') + f'"{warmup}" is not a supported type for warming up, valid' ' types are "constant" and "linear"' + ) if warmup is not None: - assert warmup_iters > 0, \ - '"warmup_iters" must be a positive integer' - assert 0 < warmup_ratio <= 1.0, \ - '"warmup_ratio" must be in range (0,1]' + assert warmup_iters > 0, '"warmup_iters" must be a positive integer' + assert 0 < warmup_ratio <= 1.0, '"warmup_ratio" must be in range (0,1]' self.by_epoch = by_epoch self.warmup = warmup @@ -59,11 +52,10 @@ def _set_lr(self, runner, lr_groups): if isinstance(runner.optimizer, dict): for k, optim in runner.optimizer.items(): for param_group, lr in zip(optim.param_groups, lr_groups[k]): - param_group['lr'] = lr + param_group["lr"] = lr else: - for param_group, lr in zip(runner.optimizer.param_groups, - lr_groups): - param_group['lr'] = lr + for param_group, lr in zip(runner.optimizer.param_groups, lr_groups): + param_group["lr"] = lr def get_lr(self, runner, base_lr): raise NotImplementedError @@ -72,10 +64,7 @@ def get_regular_lr(self, runner): if isinstance(runner.optimizer, dict): lr_groups = {} for k in runner.optimizer.keys(): - _lr_group = [ - self.get_lr(runner, _base_lr) - for _base_lr in self.base_lr[k] - ] + _lr_group = [self.get_lr(runner, _base_lr) for _base_lr in self.base_lr[k]] lr_groups.update({k: _lr_group}) return lr_groups @@ -83,16 +72,14 @@ def get_regular_lr(self, runner): return [self.get_lr(runner, _base_lr) for _base_lr in self.base_lr] def get_warmup_lr(self, cur_iters): - def _get_warmup_lr(cur_iters, regular_lr): - if self.warmup == 'constant': + if self.warmup == "constant": warmup_lr = [_lr * self.warmup_ratio for _lr in regular_lr] - elif self.warmup == 'linear': - k = (1 - cur_iters / self.warmup_iters) * (1 - - self.warmup_ratio) + elif self.warmup == "linear": + k = (1 - cur_iters / self.warmup_iters) * (1 - self.warmup_ratio) warmup_lr = [_lr * (1 - k) for _lr in regular_lr] - elif self.warmup == 'exp': - k = self.warmup_ratio**(1 - cur_iters / self.warmup_iters) + elif self.warmup == "exp": + k = self.warmup_ratio ** (1 - cur_iters / self.warmup_iters) warmup_lr = [_lr * k for _lr in regular_lr] return warmup_lr @@ -111,17 +98,13 @@ def before_run(self, runner): self.base_lr = {} for k, optim in runner.optimizer.items(): for group in optim.param_groups: - group.setdefault('initial_lr', group['lr']) - _base_lr = [ - group['initial_lr'] for group in optim.param_groups - ] + group.setdefault("initial_lr", group["lr"]) + _base_lr = [group["initial_lr"] for group in optim.param_groups] self.base_lr.update({k: _base_lr}) else: for group in runner.optimizer.param_groups: - group.setdefault('initial_lr', group['lr']) - self.base_lr = [ - group['initial_lr'] for group in runner.optimizer.param_groups - ] + group.setdefault("initial_lr", group["lr"]) + self.base_lr = [group["initial_lr"] for group in runner.optimizer.param_groups] def before_train_epoch(self, runner): if self.warmup_iters is None: @@ -155,7 +138,6 @@ def before_train_iter(self, runner): @HOOKS.register_module() class FixedLrUpdaterHook(LrUpdaterHook): - def __init__(self, **kwargs): super(FixedLrUpdaterHook, self).__init__(**kwargs) @@ -212,7 +194,6 @@ def get_lr(self, runner, base_lr): @HOOKS.register_module() class ExpLrUpdaterHook(LrUpdaterHook): - def __init__(self, gamma, **kwargs): self.gamma = gamma super(ExpLrUpdaterHook, self).__init__(**kwargs) @@ -224,8 +205,7 @@ def get_lr(self, runner, base_lr): @HOOKS.register_module() class PolyLrUpdaterHook(LrUpdaterHook): - - def __init__(self, power=1., min_lr=0., **kwargs): + def __init__(self, power=1.0, min_lr=0.0, **kwargs): self.power = power self.min_lr = min_lr super(PolyLrUpdaterHook, self).__init__(**kwargs) @@ -237,26 +217,24 @@ def get_lr(self, runner, base_lr): else: progress = runner.iter max_progress = runner.max_iters - coeff = (1 - progress / max_progress)**self.power + coeff = (1 - progress / max_progress) ** self.power return (base_lr - self.min_lr) * coeff + self.min_lr @HOOKS.register_module() class InvLrUpdaterHook(LrUpdaterHook): - - def __init__(self, gamma, power=1., **kwargs): + def __init__(self, gamma, power=1.0, **kwargs): self.gamma = gamma self.power = power super(InvLrUpdaterHook, self).__init__(**kwargs) def get_lr(self, runner, base_lr): progress = runner.epoch if self.by_epoch else runner.iter - return base_lr * (1 + self.gamma * progress)**(-self.power) + return base_lr * (1 + self.gamma * progress) ** (-self.power) @HOOKS.register_module() class CosineAnnealingLrUpdaterHook(LrUpdaterHook): - def __init__(self, min_lr=None, min_lr_ratio=None, **kwargs): assert (min_lr is None) ^ (min_lr_ratio is None) self.min_lr = min_lr @@ -295,17 +273,10 @@ class FlatCosineAnnealingLrUpdaterHook(LrUpdaterHook): Default: None. """ - def __init__(self, - start_percent=0.75, - min_lr=None, - min_lr_ratio=None, - **kwargs): + def __init__(self, start_percent=0.75, min_lr=None, min_lr_ratio=None, **kwargs): assert (min_lr is None) ^ (min_lr_ratio is None) - if start_percent < 0 or start_percent > 1 or not isinstance( - start_percent, float): - raise ValueError( - 'expected float between 0 and 1 start_percent, but ' - f'got {start_percent}') + if start_percent < 0 or start_percent > 1 or not isinstance(start_percent, float): + raise ValueError("expected float between 0 and 1 start_percent, but " f"got {start_percent}") self.start_percent = start_percent self.min_lr = min_lr self.min_lr_ratio = min_lr_ratio @@ -346,24 +317,18 @@ class CosineRestartLrUpdaterHook(LrUpdaterHook): Default: None. """ - def __init__(self, - periods, - restart_weights=[1], - min_lr=None, - min_lr_ratio=None, - **kwargs): + def __init__(self, periods, restart_weights=[1], min_lr=None, min_lr_ratio=None, **kwargs): assert (min_lr is None) ^ (min_lr_ratio is None) self.periods = periods self.min_lr = min_lr self.min_lr_ratio = min_lr_ratio self.restart_weights = restart_weights - assert (len(self.periods) == len(self.restart_weights) - ), 'periods and restart_weights should have the same length.' + assert len(self.periods) == len( + self.restart_weights + ), "periods and restart_weights should have the same length." super(CosineRestartLrUpdaterHook, self).__init__(**kwargs) - self.cumulative_periods = [ - sum(self.periods[0:i + 1]) for i in range(0, len(self.periods)) - ] + self.cumulative_periods = [sum(self.periods[0 : i + 1]) for i in range(0, len(self.periods))] def get_lr(self, runner, base_lr): if self.by_epoch: @@ -404,8 +369,7 @@ def get_position_from_periods(iteration, cumulative_periods): for i, period in enumerate(cumulative_periods): if iteration < period: return i - raise ValueError(f'Current iteration {iteration} exceeds ' - f'cumulative_periods {cumulative_periods}') + raise ValueError(f"Current iteration {iteration} exceeds " f"cumulative_periods {cumulative_periods}") @HOOKS.register_module() @@ -431,42 +395,38 @@ class CyclicLrUpdaterHook(LrUpdaterHook): 'linear' for linear annealing. Default: 'cos'. """ - def __init__(self, - by_epoch=False, - target_ratio=(10, 1e-4), - cyclic_times=1, - step_ratio_up=0.4, - anneal_strategy='cos', - **kwargs): + def __init__( + self, + by_epoch=False, + target_ratio=(10, 1e-4), + cyclic_times=1, + step_ratio_up=0.4, + anneal_strategy="cos", + **kwargs, + ): if isinstance(target_ratio, float): target_ratio = (target_ratio, target_ratio / 1e5) elif isinstance(target_ratio, tuple): - target_ratio = (target_ratio[0], target_ratio[0] / 1e5) \ - if len(target_ratio) == 1 else target_ratio + target_ratio = (target_ratio[0], target_ratio[0] / 1e5) if len(target_ratio) == 1 else target_ratio else: - raise ValueError('target_ratio should be either float ' - f'or tuple, got {type(target_ratio)}') + raise ValueError("target_ratio should be either float " f"or tuple, got {type(target_ratio)}") - assert len(target_ratio) == 2, \ - '"target_ratio" must be list or tuple of two floats' - assert 0 <= step_ratio_up < 1.0, \ - '"step_ratio_up" must be in range [0,1)' + assert len(target_ratio) == 2, '"target_ratio" must be list or tuple of two floats' + assert 0 <= step_ratio_up < 1.0, '"step_ratio_up" must be in range [0,1)' self.target_ratio = target_ratio self.cyclic_times = cyclic_times self.step_ratio_up = step_ratio_up self.lr_phases = [] # init lr_phases # validate anneal_strategy - if anneal_strategy not in ['cos', 'linear']: - raise ValueError('anneal_strategy must be one of "cos" or ' - f'"linear", instead got {anneal_strategy}') - elif anneal_strategy == 'cos': + if anneal_strategy not in ["cos", "linear"]: + raise ValueError('anneal_strategy must be one of "cos" or ' f'"linear", instead got {anneal_strategy}') + elif anneal_strategy == "cos": self.anneal_func = annealing_cos - elif anneal_strategy == 'linear': + elif anneal_strategy == "linear": self.anneal_func = annealing_linear - assert not by_epoch, \ - 'currently only support "by_epoch" = False' + assert not by_epoch, 'currently only support "by_epoch" = False' super(CyclicLrUpdaterHook, self).__init__(by_epoch, **kwargs) def before_run(self, runner): @@ -475,23 +435,18 @@ def before_run(self, runner): # total lr_phases are separated as up and down max_iter_per_phase = runner.max_iters // self.cyclic_times iter_up_phase = int(self.step_ratio_up * max_iter_per_phase) + self.lr_phases.append([0, iter_up_phase, max_iter_per_phase, 1, self.target_ratio[0]]) self.lr_phases.append( - [0, iter_up_phase, max_iter_per_phase, 1, self.target_ratio[0]]) - self.lr_phases.append([ - iter_up_phase, max_iter_per_phase, max_iter_per_phase, - self.target_ratio[0], self.target_ratio[1] - ]) + [iter_up_phase, max_iter_per_phase, max_iter_per_phase, self.target_ratio[0], self.target_ratio[1]] + ) def get_lr(self, runner, base_lr): curr_iter = runner.iter - for (start_iter, end_iter, max_iter_per_phase, start_ratio, - end_ratio) in self.lr_phases: + for start_iter, end_iter, max_iter_per_phase, start_ratio, end_ratio in self.lr_phases: curr_iter %= max_iter_per_phase if start_iter <= curr_iter < end_iter: progress = curr_iter - start_iter - return self.anneal_func(base_lr * start_ratio, - base_lr * end_ratio, - progress / (end_iter - start_iter)) + return self.anneal_func(base_lr * start_ratio, base_lr * end_ratio, progress / (end_iter - start_iter)) @HOOKS.register_module() @@ -529,42 +484,39 @@ class OneCycleLrUpdaterHook(LrUpdaterHook): Default: False """ - def __init__(self, - max_lr, - total_steps=None, - pct_start=0.3, - anneal_strategy='cos', - div_factor=25, - final_div_factor=1e4, - three_phase=False, - **kwargs): + def __init__( + self, + max_lr, + total_steps=None, + pct_start=0.3, + anneal_strategy="cos", + div_factor=25, + final_div_factor=1e4, + three_phase=False, + **kwargs, + ): # validate by_epoch, currently only support by_epoch = False - if 'by_epoch' not in kwargs: - kwargs['by_epoch'] = False + if "by_epoch" not in kwargs: + kwargs["by_epoch"] = False else: - assert not kwargs['by_epoch'], \ - 'currently only support "by_epoch" = False' + assert not kwargs["by_epoch"], 'currently only support "by_epoch" = False' if not isinstance(max_lr, (numbers.Number, list, dict)): - raise ValueError('the type of max_lr must be the one of list or ' - f'dict, but got {type(max_lr)}') + raise ValueError("the type of max_lr must be the one of list or " f"dict, but got {type(max_lr)}") self._max_lr = max_lr if total_steps is not None: if not isinstance(total_steps, int): - raise ValueError('the type of total_steps must be int, but' - f'got {type(total_steps)}') + raise ValueError("the type of total_steps must be int, but" f"got {type(total_steps)}") self.total_steps = total_steps # validate pct_start if pct_start < 0 or pct_start > 1 or not isinstance(pct_start, float): - raise ValueError('expected float between 0 and 1 pct_start, but ' - f'got {pct_start}') + raise ValueError("expected float between 0 and 1 pct_start, but " f"got {pct_start}") self.pct_start = pct_start # validate anneal_strategy - if anneal_strategy not in ['cos', 'linear']: - raise ValueError('anneal_strategy must be one of "cos" or ' - f'"linear", instead got {anneal_strategy}') - elif anneal_strategy == 'cos': + if anneal_strategy not in ["cos", "linear"]: + raise ValueError('anneal_strategy must be one of "cos" or ' f'"linear", instead got {anneal_strategy}') + elif anneal_strategy == "cos": self.anneal_func = annealing_cos - elif anneal_strategy == 'linear': + elif anneal_strategy == "linear": self.anneal_func = annealing_linear self.div_factor = div_factor self.final_div_factor = final_div_factor @@ -573,15 +525,16 @@ def __init__(self, super(OneCycleLrUpdaterHook, self).__init__(**kwargs) def before_run(self, runner): - if hasattr(self, 'total_steps'): + if hasattr(self, "total_steps"): total_steps = self.total_steps else: total_steps = runner.max_iters if total_steps < runner.max_iters: raise ValueError( - 'The total steps must be greater than or equal to max ' - f'iterations {runner.max_iters} of runner, but total steps ' - f'is {total_steps}.') + "The total steps must be greater than or equal to max " + f"iterations {runner.max_iters} of runner, but total steps " + f"is {total_steps}." + ) if isinstance(runner.optimizer, dict): self.base_lr = {} @@ -589,27 +542,21 @@ def before_run(self, runner): _max_lr = format_param(k, optim, self._max_lr) self.base_lr[k] = [lr / self.div_factor for lr in _max_lr] for group, lr in zip(optim.param_groups, self.base_lr[k]): - group.setdefault('initial_lr', lr) + group.setdefault("initial_lr", lr) else: k = type(runner.optimizer).__name__ _max_lr = format_param(k, runner.optimizer, self._max_lr) self.base_lr = [lr / self.div_factor for lr in _max_lr] for group, lr in zip(runner.optimizer.param_groups, self.base_lr): - group.setdefault('initial_lr', lr) + group.setdefault("initial_lr", lr) if self.three_phase: - self.lr_phases.append( - [float(self.pct_start * total_steps) - 1, 1, self.div_factor]) - self.lr_phases.append([ - float(2 * self.pct_start * total_steps) - 2, self.div_factor, 1 - ]) - self.lr_phases.append( - [total_steps - 1, 1, 1 / self.final_div_factor]) + self.lr_phases.append([float(self.pct_start * total_steps) - 1, 1, self.div_factor]) + self.lr_phases.append([float(2 * self.pct_start * total_steps) - 2, self.div_factor, 1]) + self.lr_phases.append([total_steps - 1, 1, 1 / self.final_div_factor]) else: - self.lr_phases.append( - [float(self.pct_start * total_steps) - 1, 1, self.div_factor]) - self.lr_phases.append( - [total_steps - 1, self.div_factor, 1 / self.final_div_factor]) + self.lr_phases.append([float(self.pct_start * total_steps) - 1, 1, self.div_factor]) + self.lr_phases.append([total_steps - 1, self.div_factor, 1 / self.final_div_factor]) def get_lr(self, runner, base_lr): curr_iter = runner.iter @@ -617,8 +564,7 @@ def get_lr(self, runner, base_lr): for i, (end_iter, start_lr, end_lr) in enumerate(self.lr_phases): if curr_iter <= end_iter: pct = (curr_iter - start_iter) / (end_iter - start_iter) - lr = self.anneal_func(base_lr * start_lr, base_lr * end_lr, - pct) + lr = self.anneal_func(base_lr * start_lr, base_lr * end_lr, pct) break start_iter = end_iter return lr @@ -661,10 +607,9 @@ def format_param(name, optim, param): return [param] * len(optim.param_groups) elif isinstance(param, (list, tuple)): # multi param groups if len(param) != len(optim.param_groups): - raise ValueError(f'expected {len(optim.param_groups)} ' - f'values for {name}, got {len(param)}') + raise ValueError(f"expected {len(optim.param_groups)} " f"values for {name}, got {len(param)}") return param else: # multi optimizers if name not in param: - raise KeyError(f'{name} is not found in {param.keys()}') + raise KeyError(f"{name} is not found in {param.keys()}") return param[name] diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/momentum_updater.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/momentum_updater.py index 171f1cf0..ca2a36b7 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/momentum_updater.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/momentum_updater.py @@ -1,27 +1,20 @@ # Copyright (c) OpenMMLab. All rights reserved. -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv from .hook import HOOKS, Hook from .lr_updater import annealing_cos, annealing_linear, format_param class MomentumUpdaterHook(Hook): - - def __init__(self, - by_epoch=True, - warmup=None, - warmup_iters=0, - warmup_ratio=0.9): + def __init__(self, by_epoch=True, warmup=None, warmup_iters=0, warmup_ratio=0.9): # validate the "warmup" argument if warmup is not None: - if warmup not in ['constant', 'linear', 'exp']: + if warmup not in ["constant", "linear", "exp"]: raise ValueError( - f'"{warmup}" is not a supported type for warming up, valid' - ' types are "constant" and "linear"') + f'"{warmup}" is not a supported type for warming up, valid' ' types are "constant" and "linear"' + ) if warmup is not None: - assert warmup_iters > 0, \ - '"warmup_iters" must be a positive integer' - assert 0 < warmup_ratio <= 1.0, \ - '"warmup_momentum" must be in range (0,1]' + assert warmup_iters > 0, '"warmup_iters" must be a positive integer' + assert 0 < warmup_ratio <= 1.0, '"warmup_momentum" must be in range (0,1]' self.by_epoch = by_epoch self.warmup = warmup @@ -29,25 +22,22 @@ def __init__(self, self.warmup_ratio = warmup_ratio self.base_momentum = [] # initial momentum for all param groups - self.regular_momentum = [ - ] # expected momentum if no warming up is performed + self.regular_momentum = [] # expected momentum if no warming up is performed def _set_momentum(self, runner, momentum_groups): if isinstance(runner.optimizer, dict): for k, optim in runner.optimizer.items(): - for param_group, mom in zip(optim.param_groups, - momentum_groups[k]): - if 'momentum' in param_group.keys(): - param_group['momentum'] = mom - elif 'betas' in param_group.keys(): - param_group['betas'] = (mom, param_group['betas'][1]) + for param_group, mom in zip(optim.param_groups, momentum_groups[k]): + if "momentum" in param_group.keys(): + param_group["momentum"] = mom + elif "betas" in param_group.keys(): + param_group["betas"] = (mom, param_group["betas"][1]) else: - for param_group, mom in zip(runner.optimizer.param_groups, - momentum_groups): - if 'momentum' in param_group.keys(): - param_group['momentum'] = mom - elif 'betas' in param_group.keys(): - param_group['betas'] = (mom, param_group['betas'][1]) + for param_group, mom in zip(runner.optimizer.param_groups, momentum_groups): + if "momentum" in param_group.keys(): + param_group["momentum"] = mom + elif "betas" in param_group.keys(): + param_group["betas"] = (mom, param_group["betas"][1]) def get_momentum(self, runner, base_momentum): raise NotImplementedError @@ -57,43 +47,29 @@ def get_regular_momentum(self, runner): momentum_groups = {} for k in runner.optimizer.keys(): _momentum_group = [ - self.get_momentum(runner, _base_momentum) - for _base_momentum in self.base_momentum[k] + self.get_momentum(runner, _base_momentum) for _base_momentum in self.base_momentum[k] ] momentum_groups.update({k: _momentum_group}) return momentum_groups else: - return [ - self.get_momentum(runner, _base_momentum) - for _base_momentum in self.base_momentum - ] + return [self.get_momentum(runner, _base_momentum) for _base_momentum in self.base_momentum] def get_warmup_momentum(self, cur_iters): - def _get_warmup_momentum(cur_iters, regular_momentum): - if self.warmup == 'constant': - warmup_momentum = [ - _momentum / self.warmup_ratio - for _momentum in self.regular_momentum - ] - elif self.warmup == 'linear': - k = (1 - cur_iters / self.warmup_iters) * (1 - - self.warmup_ratio) - warmup_momentum = [ - _momentum / (1 - k) for _momentum in self.regular_mom - ] - elif self.warmup == 'exp': - k = self.warmup_ratio**(1 - cur_iters / self.warmup_iters) - warmup_momentum = [ - _momentum / k for _momentum in self.regular_mom - ] + if self.warmup == "constant": + warmup_momentum = [_momentum / self.warmup_ratio for _momentum in self.regular_momentum] + elif self.warmup == "linear": + k = (1 - cur_iters / self.warmup_iters) * (1 - self.warmup_ratio) + warmup_momentum = [_momentum / (1 - k) for _momentum in self.regular_mom] + elif self.warmup == "exp": + k = self.warmup_ratio ** (1 - cur_iters / self.warmup_iters) + warmup_momentum = [_momentum / k for _momentum in self.regular_mom] return warmup_momentum if isinstance(self.regular_momentum, dict): momentum_groups = {} for key, regular_momentum in self.regular_momentum.items(): - momentum_groups[key] = _get_warmup_momentum( - cur_iters, regular_momentum) + momentum_groups[key] = _get_warmup_momentum(cur_iters, regular_momentum) return momentum_groups else: return _get_warmup_momentum(cur_iters, self.regular_momentum) @@ -106,24 +82,19 @@ def before_run(self, runner): self.base_momentum = {} for k, optim in runner.optimizer.items(): for group in optim.param_groups: - if 'momentum' in group.keys(): - group.setdefault('initial_momentum', group['momentum']) + if "momentum" in group.keys(): + group.setdefault("initial_momentum", group["momentum"]) else: - group.setdefault('initial_momentum', group['betas'][0]) - _base_momentum = [ - group['initial_momentum'] for group in optim.param_groups - ] + group.setdefault("initial_momentum", group["betas"][0]) + _base_momentum = [group["initial_momentum"] for group in optim.param_groups] self.base_momentum.update({k: _base_momentum}) else: for group in runner.optimizer.param_groups: - if 'momentum' in group.keys(): - group.setdefault('initial_momentum', group['momentum']) + if "momentum" in group.keys(): + group.setdefault("initial_momentum", group["momentum"]) else: - group.setdefault('initial_momentum', group['betas'][0]) - self.base_momentum = [ - group['initial_momentum'] - for group in runner.optimizer.param_groups - ] + group.setdefault("initial_momentum", group["betas"][0]) + self.base_momentum = [group["initial_momentum"] for group in runner.optimizer.param_groups] def before_train_epoch(self, runner): if not self.by_epoch: @@ -200,7 +171,6 @@ def get_momentum(self, runner, base_momentum): @HOOKS.register_module() class CosineAnnealingMomentumUpdaterHook(MomentumUpdaterHook): - def __init__(self, min_momentum=None, min_momentum_ratio=None, **kwargs): assert (min_momentum is None) ^ (min_momentum_ratio is None) self.min_momentum = min_momentum @@ -218,8 +188,7 @@ def get_momentum(self, runner, base_momentum): target_momentum = base_momentum * self.min_momentum_ratio else: target_momentum = self.min_momentum - return annealing_cos(base_momentum, target_momentum, - progress / max_progress) + return annealing_cos(base_momentum, target_momentum, progress / max_progress) @HOOKS.register_module() @@ -241,33 +210,23 @@ class CyclicMomentumUpdaterHook(MomentumUpdaterHook): by_epoch (bool): Whether to update momentum by epoch. """ - def __init__(self, - by_epoch=False, - target_ratio=(0.85 / 0.95, 1), - cyclic_times=1, - step_ratio_up=0.4, - **kwargs): + def __init__(self, by_epoch=False, target_ratio=(0.85 / 0.95, 1), cyclic_times=1, step_ratio_up=0.4, **kwargs): if isinstance(target_ratio, float): target_ratio = (target_ratio, target_ratio / 1e5) elif isinstance(target_ratio, tuple): - target_ratio = (target_ratio[0], target_ratio[0] / 1e5) \ - if len(target_ratio) == 1 else target_ratio + target_ratio = (target_ratio[0], target_ratio[0] / 1e5) if len(target_ratio) == 1 else target_ratio else: - raise ValueError('target_ratio should be either float ' - f'or tuple, got {type(target_ratio)}') + raise ValueError("target_ratio should be either float " f"or tuple, got {type(target_ratio)}") - assert len(target_ratio) == 2, \ - '"target_ratio" must be list or tuple of two floats' - assert 0 <= step_ratio_up < 1.0, \ - '"step_ratio_up" must be in range [0,1)' + assert len(target_ratio) == 2, '"target_ratio" must be list or tuple of two floats' + assert 0 <= step_ratio_up < 1.0, '"step_ratio_up" must be in range [0,1)' self.target_ratio = target_ratio self.cyclic_times = cyclic_times self.step_ratio_up = step_ratio_up self.momentum_phases = [] # init momentum_phases # currently only support by_epoch=False - assert not by_epoch, \ - 'currently only support "by_epoch" = False' + assert not by_epoch, 'currently only support "by_epoch" = False' super(CyclicMomentumUpdaterHook, self).__init__(by_epoch, **kwargs) def before_run(self, runner): @@ -276,23 +235,20 @@ def before_run(self, runner): # total momentum_phases are separated as up and down max_iter_per_phase = runner.max_iters // self.cyclic_times iter_up_phase = int(self.step_ratio_up * max_iter_per_phase) + self.momentum_phases.append([0, iter_up_phase, max_iter_per_phase, 1, self.target_ratio[0]]) self.momentum_phases.append( - [0, iter_up_phase, max_iter_per_phase, 1, self.target_ratio[0]]) - self.momentum_phases.append([ - iter_up_phase, max_iter_per_phase, max_iter_per_phase, - self.target_ratio[0], self.target_ratio[1] - ]) + [iter_up_phase, max_iter_per_phase, max_iter_per_phase, self.target_ratio[0], self.target_ratio[1]] + ) def get_momentum(self, runner, base_momentum): curr_iter = runner.iter - for (start_iter, end_iter, max_iter_per_phase, start_ratio, - end_ratio) in self.momentum_phases: + for start_iter, end_iter, max_iter_per_phase, start_ratio, end_ratio in self.momentum_phases: curr_iter %= max_iter_per_phase if start_iter <= curr_iter < end_iter: progress = curr_iter - start_iter - return annealing_cos(base_momentum * start_ratio, - base_momentum * end_ratio, - progress / (end_iter - start_iter)) + return annealing_cos( + base_momentum * start_ratio, base_momentum * end_ratio, progress / (end_iter - start_iter) + ) @HOOKS.register_module() @@ -330,39 +286,30 @@ class OneCycleMomentumUpdaterHook(MomentumUpdaterHook): Default: False """ - def __init__(self, - base_momentum=0.85, - max_momentum=0.95, - pct_start=0.3, - anneal_strategy='cos', - three_phase=False, - **kwargs): + def __init__( + self, base_momentum=0.85, max_momentum=0.95, pct_start=0.3, anneal_strategy="cos", three_phase=False, **kwargs + ): # validate by_epoch, currently only support by_epoch=False - if 'by_epoch' not in kwargs: - kwargs['by_epoch'] = False + if "by_epoch" not in kwargs: + kwargs["by_epoch"] = False else: - assert not kwargs['by_epoch'], \ - 'currently only support "by_epoch" = False' + assert not kwargs["by_epoch"], 'currently only support "by_epoch" = False' if not isinstance(base_momentum, (float, list, dict)): - raise ValueError('base_momentum must be the type among of float,' - 'list or dict.') + raise ValueError("base_momentum must be the type among of float," "list or dict.") self._base_momentum = base_momentum if not isinstance(max_momentum, (float, list, dict)): - raise ValueError('max_momentum must be the type among of float,' - 'list or dict.') + raise ValueError("max_momentum must be the type among of float," "list or dict.") self._max_momentum = max_momentum # validate pct_start if pct_start < 0 or pct_start > 1 or not isinstance(pct_start, float): - raise ValueError('Expected float between 0 and 1 pct_start, but ' - f'got {pct_start}') + raise ValueError("Expected float between 0 and 1 pct_start, but " f"got {pct_start}") self.pct_start = pct_start # validate anneal_strategy - if anneal_strategy not in ['cos', 'linear']: - raise ValueError('anneal_strategy must by one of "cos" or ' - f'"linear", instead got {anneal_strategy}') - elif anneal_strategy == 'cos': + if anneal_strategy not in ["cos", "linear"]: + raise ValueError('anneal_strategy must by one of "cos" or ' f'"linear", instead got {anneal_strategy}') + elif anneal_strategy == "cos": self.anneal_func = annealing_cos - elif anneal_strategy == 'linear': + elif anneal_strategy == "linear": self.anneal_func = annealing_linear self.three_phase = three_phase self.momentum_phases = [] # init momentum_phases @@ -371,107 +318,91 @@ def __init__(self, def before_run(self, runner): if isinstance(runner.optimizer, dict): for k, optim in runner.optimizer.items(): - if ('momentum' not in optim.defaults - and 'betas' not in optim.defaults): - raise ValueError('optimizer must support momentum with' - 'option enabled') - self.use_beta1 = 'betas' in optim.defaults + if "momentum" not in optim.defaults and "betas" not in optim.defaults: + raise ValueError("optimizer must support momentum with" "option enabled") + self.use_beta1 = "betas" in optim.defaults _base_momentum = format_param(k, optim, self._base_momentum) _max_momentum = format_param(k, optim, self._max_momentum) - for group, b_momentum, m_momentum in zip( - optim.param_groups, _base_momentum, _max_momentum): + for group, b_momentum, m_momentum in zip(optim.param_groups, _base_momentum, _max_momentum): if self.use_beta1: - _, beta2 = group['betas'] - group['betas'] = (m_momentum, beta2) + _, beta2 = group["betas"] + group["betas"] = (m_momentum, beta2) else: - group['momentum'] = m_momentum - group['base_momentum'] = b_momentum - group['max_momentum'] = m_momentum + group["momentum"] = m_momentum + group["base_momentum"] = b_momentum + group["max_momentum"] = m_momentum else: optim = runner.optimizer - if ('momentum' not in optim.defaults - and 'betas' not in optim.defaults): - raise ValueError('optimizer must support momentum with' - 'option enabled') - self.use_beta1 = 'betas' in optim.defaults + if "momentum" not in optim.defaults and "betas" not in optim.defaults: + raise ValueError("optimizer must support momentum with" "option enabled") + self.use_beta1 = "betas" in optim.defaults k = type(optim).__name__ _base_momentum = format_param(k, optim, self._base_momentum) _max_momentum = format_param(k, optim, self._max_momentum) - for group, b_momentum, m_momentum in zip(optim.param_groups, - _base_momentum, - _max_momentum): + for group, b_momentum, m_momentum in zip(optim.param_groups, _base_momentum, _max_momentum): if self.use_beta1: - _, beta2 = group['betas'] - group['betas'] = (m_momentum, beta2) + _, beta2 = group["betas"] + group["betas"] = (m_momentum, beta2) else: - group['momentum'] = m_momentum - group['base_momentum'] = b_momentum - group['max_momentum'] = m_momentum + group["momentum"] = m_momentum + group["base_momentum"] = b_momentum + group["max_momentum"] = m_momentum if self.three_phase: - self.momentum_phases.append({ - 'end_iter': - float(self.pct_start * runner.max_iters) - 1, - 'start_momentum': - 'max_momentum', - 'end_momentum': - 'base_momentum' - }) - self.momentum_phases.append({ - 'end_iter': - float(2 * self.pct_start * runner.max_iters) - 2, - 'start_momentum': - 'base_momentum', - 'end_momentum': - 'max_momentum' - }) - self.momentum_phases.append({ - 'end_iter': runner.max_iters - 1, - 'start_momentum': 'max_momentum', - 'end_momentum': 'max_momentum' - }) + self.momentum_phases.append( + { + "end_iter": float(self.pct_start * runner.max_iters) - 1, + "start_momentum": "max_momentum", + "end_momentum": "base_momentum", + } + ) + self.momentum_phases.append( + { + "end_iter": float(2 * self.pct_start * runner.max_iters) - 2, + "start_momentum": "base_momentum", + "end_momentum": "max_momentum", + } + ) + self.momentum_phases.append( + {"end_iter": runner.max_iters - 1, "start_momentum": "max_momentum", "end_momentum": "max_momentum"} + ) else: - self.momentum_phases.append({ - 'end_iter': - float(self.pct_start * runner.max_iters) - 1, - 'start_momentum': - 'max_momentum', - 'end_momentum': - 'base_momentum' - }) - self.momentum_phases.append({ - 'end_iter': runner.max_iters - 1, - 'start_momentum': 'base_momentum', - 'end_momentum': 'max_momentum' - }) + self.momentum_phases.append( + { + "end_iter": float(self.pct_start * runner.max_iters) - 1, + "start_momentum": "max_momentum", + "end_momentum": "base_momentum", + } + ) + self.momentum_phases.append( + {"end_iter": runner.max_iters - 1, "start_momentum": "base_momentum", "end_momentum": "max_momentum"} + ) def _set_momentum(self, runner, momentum_groups): if isinstance(runner.optimizer, dict): for k, optim in runner.optimizer.items(): - for param_group, mom in zip(optim.param_groups, - momentum_groups[k]): - if 'momentum' in param_group.keys(): - param_group['momentum'] = mom - elif 'betas' in param_group.keys(): - param_group['betas'] = (mom, param_group['betas'][1]) + for param_group, mom in zip(optim.param_groups, momentum_groups[k]): + if "momentum" in param_group.keys(): + param_group["momentum"] = mom + elif "betas" in param_group.keys(): + param_group["betas"] = (mom, param_group["betas"][1]) else: - for param_group, mom in zip(runner.optimizer.param_groups, - momentum_groups): - if 'momentum' in param_group.keys(): - param_group['momentum'] = mom - elif 'betas' in param_group.keys(): - param_group['betas'] = (mom, param_group['betas'][1]) + for param_group, mom in zip(runner.optimizer.param_groups, momentum_groups): + if "momentum" in param_group.keys(): + param_group["momentum"] = mom + elif "betas" in param_group.keys(): + param_group["betas"] = (mom, param_group["betas"][1]) def get_momentum(self, runner, param_group): curr_iter = runner.iter start_iter = 0 for i, phase in enumerate(self.momentum_phases): - end_iter = phase['end_iter'] + end_iter = phase["end_iter"] if curr_iter <= end_iter or i == len(self.momentum_phases) - 1: pct = (curr_iter - start_iter) / (end_iter - start_iter) momentum = self.anneal_func( - param_group[phase['start_momentum']], - param_group[phase['end_momentum']], pct) + param_group[phase["start_momentum"]], param_group[phase["end_momentum"]], pct + ) break start_iter = end_iter return momentum @@ -480,10 +411,7 @@ def get_regular_momentum(self, runner): if isinstance(runner.optimizer, dict): momentum_groups = {} for k, optim in runner.optimizer.items(): - _momentum_group = [ - self.get_momentum(runner, param_group) - for param_group in optim.param_groups - ] + _momentum_group = [self.get_momentum(runner, param_group) for param_group in optim.param_groups] momentum_groups.update({k: _momentum_group}) return momentum_groups else: diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/optimizer.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/optimizer.py index 6cf2610b..83a60ed7 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/optimizer.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/hooks/optimizer.py @@ -5,7 +5,7 @@ from torch.nn.utils import clip_grad -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, _BatchNorm, digit_version +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import TORCH_VERSION, _BatchNorm, digit_version from ..dist_utils import allreduce_grads from ..fp16_utils import LossScaler, wrap_fp16_model from .hook import HOOKS, Hook @@ -20,25 +20,22 @@ @HOOKS.register_module() class OptimizerHook(Hook): - def __init__(self, grad_clip=None): self.grad_clip = grad_clip def clip_grads(self, params): - params = list( - filter(lambda p: p.requires_grad and p.grad is not None, params)) + params = list(filter(lambda p: p.requires_grad and p.grad is not None, params)) if len(params) > 0: return clip_grad.clip_grad_norm_(params, **self.grad_clip) def after_train_iter(self, runner): runner.optimizer.zero_grad() - runner.outputs['loss'].backward() + runner.outputs["loss"].backward() if self.grad_clip is not None: grad_norm = self.clip_grads(runner.model.parameters()) if grad_norm is not None: # Add grad norm to the logger - runner.log_buffer.update({'grad_norm': float(grad_norm)}, - runner.outputs['num_samples']) + runner.log_buffer.update({"grad_norm": float(grad_norm)}, runner.outputs["num_samples"]) runner.optimizer.step() @@ -64,9 +61,9 @@ class GradientCumulativeOptimizerHook(OptimizerHook): def __init__(self, cumulative_iters=1, **kwargs): super(GradientCumulativeOptimizerHook, self).__init__(**kwargs) - assert isinstance(cumulative_iters, int) and cumulative_iters > 0, \ - f'cumulative_iters only accepts positive int, but got ' \ - f'{type(cumulative_iters)} instead.' + assert isinstance(cumulative_iters, int) and cumulative_iters > 0, ( + f"cumulative_iters only accepts positive int, but got " f"{type(cumulative_iters)} instead." + ) self.cumulative_iters = cumulative_iters self.divisible_iters = 0 @@ -84,20 +81,20 @@ def has_batch_norm(self, module): def _init(self, runner): if runner.iter % self.cumulative_iters != 0: runner.logger.warning( - 'Resume iter number is not divisible by cumulative_iters in ' - 'GradientCumulativeOptimizerHook, which means the gradient of ' - 'some iters is lost and the result may be influenced slightly.' + "Resume iter number is not divisible by cumulative_iters in " + "GradientCumulativeOptimizerHook, which means the gradient of " + "some iters is lost and the result may be influenced slightly." ) if self.has_batch_norm(runner.model) and self.cumulative_iters > 1: runner.logger.warning( - 'GradientCumulativeOptimizerHook may slightly decrease ' - 'performance if the model has BatchNorm layers.') + "GradientCumulativeOptimizerHook may slightly decrease " + "performance if the model has BatchNorm layers." + ) residual_iters = runner.max_iters - runner.iter - self.divisible_iters = ( - residual_iters // self.cumulative_iters * self.cumulative_iters) + self.divisible_iters = residual_iters // self.cumulative_iters * self.cumulative_iters self.remainder_iters = residual_iters - self.divisible_iters self.initialized = True @@ -110,25 +107,21 @@ def after_train_iter(self, runner): loss_factor = self.cumulative_iters else: loss_factor = self.remainder_iters - loss = runner.outputs['loss'] + loss = runner.outputs["loss"] loss = loss / loss_factor loss.backward() - if (self.every_n_iters(runner, self.cumulative_iters) - or self.is_last_iter(runner)): - + if self.every_n_iters(runner, self.cumulative_iters) or self.is_last_iter(runner): if self.grad_clip is not None: grad_norm = self.clip_grads(runner.model.parameters()) if grad_norm is not None: # Add grad norm to the logger - runner.log_buffer.update({'grad_norm': float(grad_norm)}, - runner.outputs['num_samples']) + runner.log_buffer.update({"grad_norm": float(grad_norm)}, runner.outputs["num_samples"]) runner.optimizer.step() runner.optimizer.zero_grad() -if (TORCH_VERSION != 'parrots' - and digit_version(TORCH_VERSION) >= digit_version('1.6.0')): +if TORCH_VERSION != "parrots" and digit_version(TORCH_VERSION) >= digit_version("1.6.0"): @HOOKS.register_module() class Fp16OptimizerHook(OptimizerHook): @@ -159,18 +152,13 @@ class Fp16OptimizerHook(OptimizerHook): >>> optimizer_hook = Fp16OptimizerHook(loss_scale=loss_scale) """ - def __init__(self, - grad_clip=None, - coalesce=True, - bucket_size_mb=-1, - loss_scale=512., - distributed=True): + def __init__(self, grad_clip=None, coalesce=True, bucket_size_mb=-1, loss_scale=512.0, distributed=True): self.grad_clip = grad_clip self.coalesce = coalesce self.bucket_size_mb = bucket_size_mb self.distributed = distributed self._scale_update_param = None - if loss_scale == 'dynamic': + if loss_scale == "dynamic": self.loss_scaler = GradScaler() elif isinstance(loss_scale, float): self._scale_update_param = loss_scale @@ -178,32 +166,28 @@ def __init__(self, elif isinstance(loss_scale, dict): self.loss_scaler = GradScaler(**loss_scale) else: - raise ValueError('loss_scale must be of type float, dict, or ' - f'"dynamic", got {loss_scale}') + raise ValueError("loss_scale must be of type float, dict, or " f'"dynamic", got {loss_scale}') def before_run(self, runner): """Preparing steps before Mixed Precision Training.""" # wrap model mode to fp16 wrap_fp16_model(runner.model) # resume from state dict - if 'fp16' in runner.meta and 'loss_scaler' in runner.meta['fp16']: - scaler_state_dict = runner.meta['fp16']['loss_scaler'] + if "fp16" in runner.meta and "loss_scaler" in runner.meta["fp16"]: + scaler_state_dict = runner.meta["fp16"]["loss_scaler"] self.loss_scaler.load_state_dict(scaler_state_dict) def copy_grads_to_fp32(self, fp16_net, fp32_weights): """Copy gradients from fp16 model to fp32 weight copy.""" - for fp32_param, fp16_param in zip(fp32_weights, - fp16_net.parameters()): + for fp32_param, fp16_param in zip(fp32_weights, fp16_net.parameters()): if fp16_param.grad is not None: if fp32_param.grad is None: - fp32_param.grad = fp32_param.data.new( - fp32_param.size()) + fp32_param.grad = fp32_param.data.new(fp32_param.size()) fp32_param.grad.copy_(fp16_param.grad) def copy_params_to_fp16(self, fp16_net, fp32_weights): """Copy updated params from fp32 weight copy to fp16 model.""" - for fp16_param, fp32_param in zip(fp16_net.parameters(), - fp32_weights): + for fp16_param, fp32_param in zip(fp16_net.parameters(), fp32_weights): fp16_param.data.copy_(fp32_param.data) def after_train_iter(self, runner): @@ -221,26 +205,23 @@ def after_train_iter(self, runner): runner.model.zero_grad() runner.optimizer.zero_grad() - self.loss_scaler.scale(runner.outputs['loss']).backward() + self.loss_scaler.scale(runner.outputs["loss"]).backward() self.loss_scaler.unscale_(runner.optimizer) # grad clip if self.grad_clip is not None: grad_norm = self.clip_grads(runner.model.parameters()) if grad_norm is not None: # Add grad norm to the logger - runner.log_buffer.update({'grad_norm': float(grad_norm)}, - runner.outputs['num_samples']) + runner.log_buffer.update({"grad_norm": float(grad_norm)}, runner.outputs["num_samples"]) # backward and update scaler self.loss_scaler.step(runner.optimizer) self.loss_scaler.update(self._scale_update_param) # save state_dict of loss_scaler - runner.meta.setdefault( - 'fp16', {})['loss_scaler'] = self.loss_scaler.state_dict() + runner.meta.setdefault("fp16", {})["loss_scaler"] = self.loss_scaler.state_dict() @HOOKS.register_module() - class GradientCumulativeFp16OptimizerHook(GradientCumulativeOptimizerHook, - Fp16OptimizerHook): + class GradientCumulativeFp16OptimizerHook(GradientCumulativeOptimizerHook, Fp16OptimizerHook): """Fp16 optimizer Hook (using PyTorch's implementation) implements multi-iters gradient cumulating. @@ -249,8 +230,7 @@ class GradientCumulativeFp16OptimizerHook(GradientCumulativeOptimizerHook, """ def __init__(self, *args, **kwargs): - super(GradientCumulativeFp16OptimizerHook, - self).__init__(*args, **kwargs) + super(GradientCumulativeFp16OptimizerHook, self).__init__(*args, **kwargs) def after_train_iter(self, runner): if not self.initialized: @@ -260,14 +240,12 @@ def after_train_iter(self, runner): loss_factor = self.cumulative_iters else: loss_factor = self.remainder_iters - loss = runner.outputs['loss'] + loss = runner.outputs["loss"] loss = loss / loss_factor self.loss_scaler.scale(loss).backward() - if (self.every_n_iters(runner, self.cumulative_iters) - or self.is_last_iter(runner)): - + if self.every_n_iters(runner, self.cumulative_iters) or self.is_last_iter(runner): # copy fp16 grads in the model to fp32 params in the optimizer self.loss_scaler.unscale_(runner.optimizer) @@ -275,17 +253,14 @@ def after_train_iter(self, runner): grad_norm = self.clip_grads(runner.model.parameters()) if grad_norm is not None: # Add grad norm to the logger - runner.log_buffer.update( - {'grad_norm': float(grad_norm)}, - runner.outputs['num_samples']) + runner.log_buffer.update({"grad_norm": float(grad_norm)}, runner.outputs["num_samples"]) # backward and update scaler self.loss_scaler.step(runner.optimizer) self.loss_scaler.update(self._scale_update_param) # save state_dict of loss_scaler - runner.meta.setdefault( - 'fp16', {})['loss_scaler'] = self.loss_scaler.state_dict() + runner.meta.setdefault("fp16", {})["loss_scaler"] = self.loss_scaler.state_dict() # clear grads runner.model.zero_grad() @@ -315,26 +290,19 @@ class Fp16OptimizerHook(OptimizerHook): Defaults to 512. """ - def __init__(self, - grad_clip=None, - coalesce=True, - bucket_size_mb=-1, - loss_scale=512., - distributed=True): + def __init__(self, grad_clip=None, coalesce=True, bucket_size_mb=-1, loss_scale=512.0, distributed=True): self.grad_clip = grad_clip self.coalesce = coalesce self.bucket_size_mb = bucket_size_mb self.distributed = distributed - if loss_scale == 'dynamic': - self.loss_scaler = LossScaler(mode='dynamic') + if loss_scale == "dynamic": + self.loss_scaler = LossScaler(mode="dynamic") elif isinstance(loss_scale, float): - self.loss_scaler = LossScaler( - init_scale=loss_scale, mode='static') + self.loss_scaler = LossScaler(init_scale=loss_scale, mode="static") elif isinstance(loss_scale, dict): self.loss_scaler = LossScaler(**loss_scale) else: - raise ValueError('loss_scale must be of type float, dict, or ' - f'"dynamic", got {loss_scale}') + raise ValueError("loss_scale must be of type float, dict, or " f'"dynamic", got {loss_scale}') def before_run(self, runner): """Preparing steps before Mixed Precision Training. @@ -344,15 +312,14 @@ def before_run(self, runner): """ # keep a copy of fp32 weights old_groups = runner.optimizer.param_groups - runner.optimizer.param_groups = copy.deepcopy( - runner.optimizer.param_groups) + runner.optimizer.param_groups = copy.deepcopy(runner.optimizer.param_groups) state = defaultdict(dict) p_map = { old_p: p for old_p, p in zip( - chain(*(g['params'] for g in old_groups)), - chain(*(g['params'] - for g in runner.optimizer.param_groups))) + chain(*(g["params"] for g in old_groups)), + chain(*(g["params"] for g in runner.optimizer.param_groups)), + ) } for k, v in runner.optimizer.state.items(): state[p_map[k]] = v @@ -360,24 +327,21 @@ def before_run(self, runner): # convert model to fp16 wrap_fp16_model(runner.model) # resume from state dict - if 'fp16' in runner.meta and 'loss_scaler' in runner.meta['fp16']: - scaler_state_dict = runner.meta['fp16']['loss_scaler'] + if "fp16" in runner.meta and "loss_scaler" in runner.meta["fp16"]: + scaler_state_dict = runner.meta["fp16"]["loss_scaler"] self.loss_scaler.load_state_dict(scaler_state_dict) def copy_grads_to_fp32(self, fp16_net, fp32_weights): """Copy gradients from fp16 model to fp32 weight copy.""" - for fp32_param, fp16_param in zip(fp32_weights, - fp16_net.parameters()): + for fp32_param, fp16_param in zip(fp32_weights, fp16_net.parameters()): if fp16_param.grad is not None: if fp32_param.grad is None: - fp32_param.grad = fp32_param.data.new( - fp32_param.size()) + fp32_param.grad = fp32_param.data.new(fp32_param.size()) fp32_param.grad.copy_(fp16_param.grad) def copy_params_to_fp16(self, fp16_net, fp32_weights): """Copy updated params from fp32 weight copy to fp16 model.""" - for fp16_param, fp32_param in zip(fp16_net.parameters(), - fp32_weights): + for fp16_param, fp32_param in zip(fp16_net.parameters(), fp32_weights): fp16_param.data.copy_(fp32_param.data) def after_train_iter(self, runner): @@ -395,18 +359,17 @@ def after_train_iter(self, runner): runner.model.zero_grad() runner.optimizer.zero_grad() # scale the loss value - scaled_loss = runner.outputs['loss'] * self.loss_scaler.loss_scale + scaled_loss = runner.outputs["loss"] * self.loss_scaler.loss_scale scaled_loss.backward() # copy fp16 grads in the model to fp32 params in the optimizer fp32_weights = [] for param_group in runner.optimizer.param_groups: - fp32_weights += param_group['params'] + fp32_weights += param_group["params"] self.copy_grads_to_fp32(runner.model, fp32_weights) # allreduce grads if self.distributed: - allreduce_grads(fp32_weights, self.coalesce, - self.bucket_size_mb) + allreduce_grads(fp32_weights, self.coalesce, self.bucket_size_mb) has_overflow = self.loss_scaler.has_overflow(fp32_weights) # if has overflow, skip this iteration @@ -419,31 +382,25 @@ def after_train_iter(self, runner): grad_norm = self.clip_grads(fp32_weights) if grad_norm is not None: # Add grad norm to the logger - runner.log_buffer.update( - {'grad_norm': float(grad_norm)}, - runner.outputs['num_samples']) + runner.log_buffer.update({"grad_norm": float(grad_norm)}, runner.outputs["num_samples"]) # update fp32 params runner.optimizer.step() # copy fp32 params to the fp16 model self.copy_params_to_fp16(runner.model, fp32_weights) self.loss_scaler.update_scale(has_overflow) if has_overflow: - runner.logger.warning('Check overflow, downscale loss scale ' - f'to {self.loss_scaler.cur_scale}') + runner.logger.warning("Check overflow, downscale loss scale " f"to {self.loss_scaler.cur_scale}") # save state_dict of loss_scaler - runner.meta.setdefault( - 'fp16', {})['loss_scaler'] = self.loss_scaler.state_dict() + runner.meta.setdefault("fp16", {})["loss_scaler"] = self.loss_scaler.state_dict() @HOOKS.register_module() - class GradientCumulativeFp16OptimizerHook(GradientCumulativeOptimizerHook, - Fp16OptimizerHook): + class GradientCumulativeFp16OptimizerHook(GradientCumulativeOptimizerHook, Fp16OptimizerHook): """Fp16 optimizer Hook (using mmcv implementation) implements multi- iters gradient cumulating.""" def __init__(self, *args, **kwargs): - super(GradientCumulativeFp16OptimizerHook, - self).__init__(*args, **kwargs) + super(GradientCumulativeFp16OptimizerHook, self).__init__(*args, **kwargs) def after_train_iter(self, runner): if not self.initialized: @@ -454,25 +411,22 @@ def after_train_iter(self, runner): else: loss_factor = self.remainder_iters - loss = runner.outputs['loss'] + loss = runner.outputs["loss"] loss = loss / loss_factor # scale the loss value scaled_loss = loss * self.loss_scaler.loss_scale scaled_loss.backward() - if (self.every_n_iters(runner, self.cumulative_iters) - or self.is_last_iter(runner)): - + if self.every_n_iters(runner, self.cumulative_iters) or self.is_last_iter(runner): # copy fp16 grads in the model to fp32 params in the optimizer fp32_weights = [] for param_group in runner.optimizer.param_groups: - fp32_weights += param_group['params'] + fp32_weights += param_group["params"] self.copy_grads_to_fp32(runner.model, fp32_weights) # allreduce grads if self.distributed: - allreduce_grads(fp32_weights, self.coalesce, - self.bucket_size_mb) + allreduce_grads(fp32_weights, self.coalesce, self.bucket_size_mb) has_overflow = self.loss_scaler.has_overflow(fp32_weights) # if has overflow, skip this iteration @@ -485,23 +439,18 @@ def after_train_iter(self, runner): grad_norm = self.clip_grads(fp32_weights) if grad_norm is not None: # Add grad norm to the logger - runner.log_buffer.update( - {'grad_norm': float(grad_norm)}, - runner.outputs['num_samples']) + runner.log_buffer.update({"grad_norm": float(grad_norm)}, runner.outputs["num_samples"]) # update fp32 params runner.optimizer.step() # copy fp32 params to the fp16 model self.copy_params_to_fp16(runner.model, fp32_weights) else: - runner.logger.warning( - 'Check overflow, downscale loss scale ' - f'to {self.loss_scaler.cur_scale}') + runner.logger.warning("Check overflow, downscale loss scale " f"to {self.loss_scaler.cur_scale}") self.loss_scaler.update_scale(has_overflow) # save state_dict of loss_scaler - runner.meta.setdefault( - 'fp16', {})['loss_scaler'] = self.loss_scaler.state_dict() + runner.meta.setdefault("fp16", {})["loss_scaler"] = self.loss_scaler.state_dict() # clear grads runner.model.zero_grad() diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/iter_based_runner.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/iter_based_runner.py index b6140cbb..94204ff4 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/iter_based_runner.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/iter_based_runner.py @@ -8,7 +8,7 @@ import torch from torch.optim import Optimizer -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv from .base_runner import BaseRunner from .builder import RUNNERS from .checkpoint import save_checkpoint @@ -17,7 +17,6 @@ class IterLoader: - def __init__(self, dataloader): self._dataloader = dataloader self.iter_loader = iter(self._dataloader) @@ -32,7 +31,7 @@ def __next__(self): data = next(self.iter_loader) except StopIteration: self._epoch += 1 - if hasattr(self._dataloader.sampler, 'set_epoch'): + if hasattr(self._dataloader.sampler, "set_epoch"): self._dataloader.sampler.set_epoch(self._epoch) time.sleep(2) # Prevent possible deadlock during epoch transition self.iter_loader = iter(self._dataloader) @@ -53,35 +52,35 @@ class IterBasedRunner(BaseRunner): def train(self, data_loader, **kwargs): self.model.train() - self.mode = 'train' + self.mode = "train" self.data_loader = data_loader self._epoch = data_loader.epoch data_batch = next(data_loader) - self.call_hook('before_train_iter') + self.call_hook("before_train_iter") outputs = self.model.train_step(data_batch, self.optimizer, **kwargs) if not isinstance(outputs, dict): - raise TypeError('model.train_step() must return a dict') - if 'log_vars' in outputs: - self.log_buffer.update(outputs['log_vars'], outputs['num_samples']) + raise TypeError("model.train_step() must return a dict") + if "log_vars" in outputs: + self.log_buffer.update(outputs["log_vars"], outputs["num_samples"]) self.outputs = outputs - self.call_hook('after_train_iter') + self.call_hook("after_train_iter") self._inner_iter += 1 self._iter += 1 @torch.no_grad() def val(self, data_loader, **kwargs): self.model.eval() - self.mode = 'val' + self.mode = "val" self.data_loader = data_loader data_batch = next(data_loader) - self.call_hook('before_val_iter') + self.call_hook("before_val_iter") outputs = self.model.val_step(data_batch, **kwargs) if not isinstance(outputs, dict): - raise TypeError('model.val_step() must return a dict') - if 'log_vars' in outputs: - self.log_buffer.update(outputs['log_vars'], outputs['num_samples']) + raise TypeError("model.val_step() must return a dict") + if "log_vars" in outputs: + self.log_buffer.update(outputs["log_vars"], outputs["num_samples"]) self.outputs = outputs - self.call_hook('after_val_iter') + self.call_hook("after_val_iter") self._inner_iter += 1 def run(self, data_loaders, workflow, max_iters=None, **kwargs): @@ -100,47 +99,38 @@ def run(self, data_loaders, workflow, max_iters=None, **kwargs): assert len(data_loaders) == len(workflow) if max_iters is not None: warnings.warn( - 'setting max_iters in run is deprecated, ' - 'please set max_iters in runner_config', DeprecationWarning) + "setting max_iters in run is deprecated, " "please set max_iters in runner_config", DeprecationWarning + ) self._max_iters = max_iters - assert self._max_iters is not None, ( - 'max_iters must be specified during instantiation') - - work_dir = self.work_dir if self.work_dir is not None else 'NONE' - self.logger.info('Start running, host: %s, work_dir: %s', - get_host_info(), work_dir) - self.logger.info('Hooks will be executed in the following order:\n%s', - self.get_hook_info()) - self.logger.info('workflow: %s, max: %d iters', workflow, - self._max_iters) - self.call_hook('before_run') + assert self._max_iters is not None, "max_iters must be specified during instantiation" + + work_dir = self.work_dir if self.work_dir is not None else "NONE" + self.logger.info("Start running, host: %s, work_dir: %s", get_host_info(), work_dir) + self.logger.info("Hooks will be executed in the following order:\n%s", self.get_hook_info()) + self.logger.info("workflow: %s, max: %d iters", workflow, self._max_iters) + self.call_hook("before_run") iter_loaders = [IterLoader(x) for x in data_loaders] - self.call_hook('before_epoch') + self.call_hook("before_epoch") while self.iter < self._max_iters: for i, flow in enumerate(workflow): self._inner_iter = 0 mode, iters = flow if not isinstance(mode, str) or not hasattr(self, mode): - raise ValueError( - 'runner has no method named "{}" to run a workflow'. - format(mode)) + raise ValueError('runner has no method named "{}" to run a workflow'.format(mode)) iter_runner = getattr(self, mode) for _ in range(iters): - if mode == 'train' and self.iter >= self._max_iters: + if mode == "train" and self.iter >= self._max_iters: break iter_runner(iter_loaders[i], **kwargs) time.sleep(1) # wait for some hooks like loggers to finish - self.call_hook('after_epoch') - self.call_hook('after_run') + self.call_hook("after_epoch") + self.call_hook("after_run") - def resume(self, - checkpoint, - resume_optimizer=True, - map_location='default'): + def resume(self, checkpoint, resume_optimizer=True, map_location="default"): """Resume model from checkpoint. Args: @@ -150,38 +140,29 @@ def resume(self, map_location (str, optional): Same as :func:`torch.load`. Default to 'default'. """ - if map_location == 'default': + if map_location == "default": device_id = torch.cuda.current_device() - checkpoint = self.load_checkpoint( - checkpoint, - map_location=lambda storage, loc: storage.cuda(device_id)) + checkpoint = self.load_checkpoint(checkpoint, map_location=lambda storage, loc: storage.cuda(device_id)) else: - checkpoint = self.load_checkpoint( - checkpoint, map_location=map_location) + checkpoint = self.load_checkpoint(checkpoint, map_location=map_location) - self._epoch = checkpoint['meta']['epoch'] - self._iter = checkpoint['meta']['iter'] - self._inner_iter = checkpoint['meta']['iter'] - if 'optimizer' in checkpoint and resume_optimizer: + self._epoch = checkpoint["meta"]["epoch"] + self._iter = checkpoint["meta"]["iter"] + self._inner_iter = checkpoint["meta"]["iter"] + if "optimizer" in checkpoint and resume_optimizer: if isinstance(self.optimizer, Optimizer): - self.optimizer.load_state_dict(checkpoint['optimizer']) + self.optimizer.load_state_dict(checkpoint["optimizer"]) elif isinstance(self.optimizer, dict): for k in self.optimizer.keys(): - self.optimizer[k].load_state_dict( - checkpoint['optimizer'][k]) + self.optimizer[k].load_state_dict(checkpoint["optimizer"][k]) else: - raise TypeError( - 'Optimizer should be dict or torch.optim.Optimizer ' - f'but got {type(self.optimizer)}') - - self.logger.info(f'resumed from epoch: {self.epoch}, iter {self.iter}') - - def save_checkpoint(self, - out_dir, - filename_tmpl='iter_{}.pth', - meta=None, - save_optimizer=True, - create_symlink=True): + raise TypeError("Optimizer should be dict or torch.optim.Optimizer " f"but got {type(self.optimizer)}") + + self.logger.info(f"resumed from epoch: {self.epoch}, iter {self.iter}") + + def save_checkpoint( + self, out_dir, filename_tmpl="iter_{}.pth", meta=None, save_optimizer=True, create_symlink=True + ): """Save checkpoint to file. Args: @@ -198,8 +179,7 @@ def save_checkpoint(self, if meta is None: meta = {} elif not isinstance(meta, dict): - raise TypeError( - f'meta should be a dict or None, but got {type(meta)}') + raise TypeError(f"meta should be a dict or None, but got {type(meta)}") if self.meta is not None: meta.update(self.meta) # Note: meta.update(self.meta) should be done before @@ -215,19 +195,21 @@ def save_checkpoint(self, # in some environments, `os.symlink` is not supported, you may need to # set `create_symlink` to False if create_symlink: - dst_file = osp.join(out_dir, 'latest.pth') - if platform.system() != 'Windows': + dst_file = osp.join(out_dir, "latest.pth") + if platform.system() != "Windows": mmcv.symlink(filename, dst_file) else: shutil.copy(filepath, dst_file) - def register_training_hooks(self, - lr_config, - optimizer_config=None, - checkpoint_config=None, - log_config=None, - momentum_config=None, - custom_hooks_config=None): + def register_training_hooks( + self, + lr_config, + optimizer_config=None, + checkpoint_config=None, + log_config=None, + momentum_config=None, + custom_hooks_config=None, + ): """Register default hooks for iter-based training. Checkpoint hook, optimizer stepper hook and logger hooks will be set to @@ -257,12 +239,12 @@ def register_training_hooks(self, will be triggered after default hooks. """ if checkpoint_config is not None: - checkpoint_config.setdefault('by_epoch', False) + checkpoint_config.setdefault("by_epoch", False) if lr_config is not None: - lr_config.setdefault('by_epoch', False) + lr_config.setdefault("by_epoch", False) if log_config is not None: - for info in log_config['hooks']: - info.setdefault('by_epoch', False) + for info in log_config["hooks"]: + info.setdefault("by_epoch", False) super(IterBasedRunner, self).register_training_hooks( lr_config=lr_config, momentum_config=momentum_config, @@ -270,4 +252,5 @@ def register_training_hooks(self, checkpoint_config=checkpoint_config, log_config=log_config, timer_config=IterTimerHook(), - custom_hooks_config=custom_hooks_config) + custom_hooks_config=custom_hooks_config, + ) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/optimizer/default_constructor.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/optimizer/default_constructor.py index d804d4b1..1df9485d 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/optimizer/default_constructor.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/optimizer/default_constructor.py @@ -4,8 +4,13 @@ import torch from torch.nn import GroupNorm, LayerNorm -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import _BatchNorm, _InstanceNorm, build_from_cfg, is_list_of -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.ext_loader import check_ops_exist +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import ( + _BatchNorm, + _InstanceNorm, + build_from_cfg, + is_list_of, +) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.ext_loader import check_ops_exist from .builder import OPTIMIZER_BUILDERS, OPTIMIZERS @@ -94,47 +99,47 @@ class DefaultOptimizerConstructor: def __init__(self, optimizer_cfg, paramwise_cfg=None): if not isinstance(optimizer_cfg, dict): - raise TypeError('optimizer_cfg should be a dict', - f'but got {type(optimizer_cfg)}') + raise TypeError("optimizer_cfg should be a dict", f"but got {type(optimizer_cfg)}") self.optimizer_cfg = optimizer_cfg self.paramwise_cfg = {} if paramwise_cfg is None else paramwise_cfg - self.base_lr = optimizer_cfg.get('lr', None) - self.base_wd = optimizer_cfg.get('weight_decay', None) + self.base_lr = optimizer_cfg.get("lr", None) + self.base_wd = optimizer_cfg.get("weight_decay", None) self._validate_cfg() def _validate_cfg(self): if not isinstance(self.paramwise_cfg, dict): - raise TypeError('paramwise_cfg should be None or a dict, ' - f'but got {type(self.paramwise_cfg)}') + raise TypeError("paramwise_cfg should be None or a dict, " f"but got {type(self.paramwise_cfg)}") - if 'custom_keys' in self.paramwise_cfg: - if not isinstance(self.paramwise_cfg['custom_keys'], dict): + if "custom_keys" in self.paramwise_cfg: + if not isinstance(self.paramwise_cfg["custom_keys"], dict): raise TypeError( - 'If specified, custom_keys must be a dict, ' - f'but got {type(self.paramwise_cfg["custom_keys"])}') + "If specified, custom_keys must be a dict, " f'but got {type(self.paramwise_cfg["custom_keys"])}' + ) if self.base_wd is None: - for key in self.paramwise_cfg['custom_keys']: - if 'decay_mult' in self.paramwise_cfg['custom_keys'][key]: - raise ValueError('base_wd should not be None') + for key in self.paramwise_cfg["custom_keys"]: + if "decay_mult" in self.paramwise_cfg["custom_keys"][key]: + raise ValueError("base_wd should not be None") # get base lr and weight decay # weight_decay must be explicitly specified if mult is specified - if ('bias_decay_mult' in self.paramwise_cfg - or 'norm_decay_mult' in self.paramwise_cfg - or 'dwconv_decay_mult' in self.paramwise_cfg): + if ( + "bias_decay_mult" in self.paramwise_cfg + or "norm_decay_mult" in self.paramwise_cfg + or "dwconv_decay_mult" in self.paramwise_cfg + ): if self.base_wd is None: - raise ValueError('base_wd should not be None') + raise ValueError("base_wd should not be None") def _is_in(self, param_group, param_group_list): assert is_list_of(param_group_list, dict) - param = set(param_group['params']) + param = set(param_group["params"]) param_set = set() for group in param_group_list: - param_set.update(set(group['params'])) + param_set.update(set(group["params"])) return not param.isdisjoint(param_set) - def add_params(self, params, module, prefix='', is_dcn_module=None): + def add_params(self, params, module, prefix="", is_dcn_module=None): """Add all parameters of module to the params list. The parameters of the given module will be added to the list of param @@ -150,100 +155,91 @@ def add_params(self, params, module, prefix='', is_dcn_module=None): control conv_offset layer's learning rate. Defaults to None. """ # get param-wise options - custom_keys = self.paramwise_cfg.get('custom_keys', {}) + custom_keys = self.paramwise_cfg.get("custom_keys", {}) # first sort with alphabet order and then sort with reversed len of str sorted_keys = sorted(sorted(custom_keys.keys()), key=len, reverse=True) - bias_lr_mult = self.paramwise_cfg.get('bias_lr_mult', 1.) - bias_decay_mult = self.paramwise_cfg.get('bias_decay_mult', 1.) - norm_decay_mult = self.paramwise_cfg.get('norm_decay_mult', 1.) - dwconv_decay_mult = self.paramwise_cfg.get('dwconv_decay_mult', 1.) - bypass_duplicate = self.paramwise_cfg.get('bypass_duplicate', False) - dcn_offset_lr_mult = self.paramwise_cfg.get('dcn_offset_lr_mult', 1.) + bias_lr_mult = self.paramwise_cfg.get("bias_lr_mult", 1.0) + bias_decay_mult = self.paramwise_cfg.get("bias_decay_mult", 1.0) + norm_decay_mult = self.paramwise_cfg.get("norm_decay_mult", 1.0) + dwconv_decay_mult = self.paramwise_cfg.get("dwconv_decay_mult", 1.0) + bypass_duplicate = self.paramwise_cfg.get("bypass_duplicate", False) + dcn_offset_lr_mult = self.paramwise_cfg.get("dcn_offset_lr_mult", 1.0) # special rules for norm layers and depth-wise conv layers - is_norm = isinstance(module, - (_BatchNorm, _InstanceNorm, GroupNorm, LayerNorm)) - is_dwconv = ( - isinstance(module, torch.nn.Conv2d) - and module.in_channels == module.groups) + is_norm = isinstance(module, (_BatchNorm, _InstanceNorm, GroupNorm, LayerNorm)) + is_dwconv = isinstance(module, torch.nn.Conv2d) and module.in_channels == module.groups for name, param in module.named_parameters(recurse=False): - param_group = {'params': [param]} + param_group = {"params": [param]} if not param.requires_grad: params.append(param_group) continue if bypass_duplicate and self._is_in(param_group, params): - warnings.warn(f'{prefix} is duplicate. It is skipped since ' - f'bypass_duplicate={bypass_duplicate}') + warnings.warn(f"{prefix} is duplicate. It is skipped since " f"bypass_duplicate={bypass_duplicate}") continue # if the parameter match one of the custom keys, ignore other rules is_custom = False for key in sorted_keys: - if key in f'{prefix}.{name}': + if key in f"{prefix}.{name}": is_custom = True - lr_mult = custom_keys[key].get('lr_mult', 1.) - param_group['lr'] = self.base_lr * lr_mult + lr_mult = custom_keys[key].get("lr_mult", 1.0) + param_group["lr"] = self.base_lr * lr_mult if self.base_wd is not None: - decay_mult = custom_keys[key].get('decay_mult', 1.) - param_group['weight_decay'] = self.base_wd * decay_mult + decay_mult = custom_keys[key].get("decay_mult", 1.0) + param_group["weight_decay"] = self.base_wd * decay_mult break if not is_custom: # bias_lr_mult affects all bias parameters # except for norm.bias dcn.conv_offset.bias - if name == 'bias' and not (is_norm or is_dcn_module): - param_group['lr'] = self.base_lr * bias_lr_mult + if name == "bias" and not (is_norm or is_dcn_module): + param_group["lr"] = self.base_lr * bias_lr_mult - if (prefix.find('conv_offset') != -1 and is_dcn_module - and isinstance(module, torch.nn.Conv2d)): + if prefix.find("conv_offset") != -1 and is_dcn_module and isinstance(module, torch.nn.Conv2d): # deal with both dcn_offset's bias & weight - param_group['lr'] = self.base_lr * dcn_offset_lr_mult + param_group["lr"] = self.base_lr * dcn_offset_lr_mult # apply weight decay policies if self.base_wd is not None: # norm decay if is_norm: - param_group[ - 'weight_decay'] = self.base_wd * norm_decay_mult + param_group["weight_decay"] = self.base_wd * norm_decay_mult # depth-wise conv elif is_dwconv: - param_group[ - 'weight_decay'] = self.base_wd * dwconv_decay_mult + param_group["weight_decay"] = self.base_wd * dwconv_decay_mult # bias lr and decay - elif name == 'bias' and not is_dcn_module: + elif name == "bias" and not is_dcn_module: # TODO: current bias_decay_mult will have affect on DCN - param_group[ - 'weight_decay'] = self.base_wd * bias_decay_mult + param_group["weight_decay"] = self.base_wd * bias_decay_mult params.append(param_group) if check_ops_exist(): - from comfy_controlnet_preprocessors.uniformer.mmcv.ops import DeformConv2d, ModulatedDeformConv2d - is_dcn_module = isinstance(module, - (DeformConv2d, ModulatedDeformConv2d)) + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.ops import ( + DeformConv2d, + ModulatedDeformConv2d, + ) + + is_dcn_module = isinstance(module, (DeformConv2d, ModulatedDeformConv2d)) else: is_dcn_module = False for child_name, child_mod in module.named_children(): - child_prefix = f'{prefix}.{child_name}' if prefix else child_name - self.add_params( - params, - child_mod, - prefix=child_prefix, - is_dcn_module=is_dcn_module) + child_prefix = f"{prefix}.{child_name}" if prefix else child_name + self.add_params(params, child_mod, prefix=child_prefix, is_dcn_module=is_dcn_module) def __call__(self, model): - if hasattr(model, 'module'): + if hasattr(model, "module"): model = model.module optimizer_cfg = self.optimizer_cfg.copy() # if no paramwise option is specified, just use the global setting if not self.paramwise_cfg: - optimizer_cfg['params'] = model.parameters() + optimizer_cfg["params"] = model.parameters() return build_from_cfg(optimizer_cfg, OPTIMIZERS) # set param-wise lr and weight decay recursively params = [] self.add_params(params, model) - optimizer_cfg['params'] = params + optimizer_cfg["params"] = params return build_from_cfg(optimizer_cfg, OPTIMIZERS) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/utils.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/utils.py index 2926b93e..8ac99a1e 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/utils.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/runner/utils.py @@ -10,7 +10,7 @@ import numpy as np import torch -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv def get_host_info(): @@ -19,17 +19,17 @@ def get_host_info(): Return empty string if exception raised, e.g. ``getpass.getuser()`` will lead to error in docker container """ - host = '' + host = "" try: - host = f'{getuser()}@{gethostname()}' + host = f"{getuser()}@{gethostname()}" except Exception as e: - warnings.warn(f'Host or user not found: {str(e)}') + warnings.warn(f"Host or user not found: {str(e)}") finally: return host def get_time_str(): - return time.strftime('%Y%m%d_%H%M%S', time.localtime()) + return time.strftime("%Y%m%d_%H%M%S", time.localtime()) def obj_from_dict(info, parent=None, default_args=None): @@ -49,18 +49,17 @@ def obj_from_dict(info, parent=None, default_args=None): Returns: any type: Object built from the dict. """ - assert isinstance(info, dict) and 'type' in info + assert isinstance(info, dict) and "type" in info assert isinstance(default_args, dict) or default_args is None args = info.copy() - obj_type = args.pop('type') + obj_type = args.pop("type") if mmcv.is_str(obj_type): if parent is not None: obj_type = getattr(parent, obj_type) else: obj_type = sys.modules[obj_type] elif not isinstance(obj_type, type): - raise TypeError('type must be a str or valid type, but ' - f'got {type(obj_type)}') + raise TypeError("type must be a str or valid type, but " f"got {type(obj_type)}") if default_args is not None: for name, value in default_args.items(): args.setdefault(name, value) @@ -87,7 +86,7 @@ def set_random_seed(seed, deterministic=False, use_rank_shift=False): torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) - os.environ['PYTHONHASHSEED'] = str(seed) + os.environ["PYTHONHASHSEED"] = str(seed) if deterministic: torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/config.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/config.py index fa33c62f..49763bc0 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/config.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/config.py @@ -18,19 +18,18 @@ from .misc import import_modules_from_strings from .path import check_file_exist -if platform.system() == 'Windows': +if platform.system() == "Windows": import regex as re else: import re -BASE_KEY = '_base_' -DELETE_KEY = '_delete_' -DEPRECATION_KEY = '_deprecation_' -RESERVED_KEYS = ['filename', 'text', 'pretty_text'] +BASE_KEY = "_base_" +DELETE_KEY = "_delete_" +DEPRECATION_KEY = "_deprecation_" +RESERVED_KEYS = ["filename", "text", "pretty_text"] class ConfigDict(Dict): - def __missing__(self, name): raise KeyError(name) @@ -38,8 +37,7 @@ def __getattr__(self, name): try: value = super(ConfigDict, self).__getattr__(name) except KeyError: - ex = AttributeError(f"'{self.__class__.__name__}' object has no " - f"attribute '{name}'") + ex = AttributeError(f"'{self.__class__.__name__}' object has no " f"attribute '{name}'") except Exception as e: ex = e else: @@ -47,22 +45,22 @@ def __getattr__(self, name): raise ex -def add_args(parser, cfg, prefix=''): +def add_args(parser, cfg, prefix=""): for k, v in cfg.items(): if isinstance(v, str): - parser.add_argument('--' + prefix + k) + parser.add_argument("--" + prefix + k) elif isinstance(v, int): - parser.add_argument('--' + prefix + k, type=int) + parser.add_argument("--" + prefix + k, type=int) elif isinstance(v, float): - parser.add_argument('--' + prefix + k, type=float) + parser.add_argument("--" + prefix + k, type=float) elif isinstance(v, bool): - parser.add_argument('--' + prefix + k, action='store_true') + parser.add_argument("--" + prefix + k, action="store_true") elif isinstance(v, dict): - add_args(parser, v, prefix + k + '.') + add_args(parser, v, prefix + k + ".") elif isinstance(v, abc.Iterable): - parser.add_argument('--' + prefix + k, type=type(v[0]), nargs='+') + parser.add_argument("--" + prefix + k, type=type(v[0]), nargs="+") else: - print(f'cannot parse key {prefix + k} of type {type(v)}') + print(f"cannot parse key {prefix + k} of type {type(v)}") return parser @@ -93,14 +91,13 @@ class Config: @staticmethod def _validate_py_syntax(filename): - with open(filename, 'r', encoding='utf-8') as f: + with open(filename, "r", encoding="utf-8") as f: # Setting encoding explicitly to resolve coding issue on windows content = f.read() try: ast.parse(content) except SyntaxError as e: - raise SyntaxError('There are syntax errors in config ' - f'file {filename}: {e}') + raise SyntaxError("There are syntax errors in config " f"file {filename}: {e}") @staticmethod def _substitute_predefined_vars(filename, temp_config_name): @@ -112,33 +109,34 @@ def _substitute_predefined_vars(filename, temp_config_name): fileDirname=file_dirname, fileBasename=file_basename, fileBasenameNoExtension=file_basename_no_extension, - fileExtname=file_extname) - with open(filename, 'r', encoding='utf-8') as f: + fileExtname=file_extname, + ) + with open(filename, "r", encoding="utf-8") as f: # Setting encoding explicitly to resolve coding issue on windows config_file = f.read() for key, value in support_templates.items(): - regexp = r'\{\{\s*' + str(key) + r'\s*\}\}' - value = value.replace('\\', '/') + regexp = r"\{\{\s*" + str(key) + r"\s*\}\}" + value = value.replace("\\", "/") config_file = re.sub(regexp, value, config_file) - with open(temp_config_name, 'w', encoding='utf-8') as tmp_config_file: + with open(temp_config_name, "w", encoding="utf-8") as tmp_config_file: tmp_config_file.write(config_file) @staticmethod def _pre_substitute_base_vars(filename, temp_config_name): """Substitute base variable placehoders to string, so that parsing would work.""" - with open(filename, 'r', encoding='utf-8') as f: + with open(filename, "r", encoding="utf-8") as f: # Setting encoding explicitly to resolve coding issue on windows config_file = f.read() base_var_dict = {} - regexp = r'\{\{\s*' + BASE_KEY + r'\.([\w\.]+)\s*\}\}' + regexp = r"\{\{\s*" + BASE_KEY + r"\.([\w\.]+)\s*\}\}" base_vars = set(re.findall(regexp, config_file)) for base_var in base_vars: - randstr = f'_{base_var}_{uuid.uuid4().hex.lower()[:6]}' + randstr = f"_{base_var}_{uuid.uuid4().hex.lower()[:6]}" base_var_dict[randstr] = base_var - regexp = r'\{\{\s*' + BASE_KEY + r'\.' + base_var + r'\s*\}\}' + regexp = r"\{\{\s*" + BASE_KEY + r"\." + base_var + r"\s*\}\}" config_file = re.sub(regexp, f'"{randstr}"', config_file) - with open(temp_config_name, 'w', encoding='utf-8') as tmp_config_file: + with open(temp_config_name, "w", encoding="utf-8") as tmp_config_file: tmp_config_file.write(config_file) return base_var_dict @@ -151,24 +149,18 @@ def _substitute_base_vars(cfg, base_var_dict, base_cfg): for k, v in cfg.items(): if isinstance(v, str) and v in base_var_dict: new_v = base_cfg - for new_k in base_var_dict[v].split('.'): + for new_k in base_var_dict[v].split("."): new_v = new_v[new_k] cfg[k] = new_v elif isinstance(v, (list, tuple, dict)): - cfg[k] = Config._substitute_base_vars( - v, base_var_dict, base_cfg) + cfg[k] = Config._substitute_base_vars(v, base_var_dict, base_cfg) elif isinstance(cfg, tuple): - cfg = tuple( - Config._substitute_base_vars(c, base_var_dict, base_cfg) - for c in cfg) + cfg = tuple(Config._substitute_base_vars(c, base_var_dict, base_cfg) for c in cfg) elif isinstance(cfg, list): - cfg = [ - Config._substitute_base_vars(c, base_var_dict, base_cfg) - for c in cfg - ] + cfg = [Config._substitute_base_vars(c, base_var_dict, base_cfg) for c in cfg] elif isinstance(cfg, str) and cfg in base_var_dict: new_v = base_cfg - for new_k in base_var_dict[cfg].split('.'): + for new_k in base_var_dict[cfg].split("."): new_v = new_v[new_k] cfg = new_v @@ -179,40 +171,34 @@ def _file2dict(filename, use_predefined_variables=True): filename = osp.abspath(osp.expanduser(filename)) check_file_exist(filename) fileExtname = osp.splitext(filename)[1] - if fileExtname not in ['.py', '.json', '.yaml', '.yml']: - raise IOError('Only py/yml/yaml/json type are supported now!') + if fileExtname not in [".py", ".json", ".yaml", ".yml"]: + raise IOError("Only py/yml/yaml/json type are supported now!") with tempfile.TemporaryDirectory() as temp_config_dir: - temp_config_file = tempfile.NamedTemporaryFile( - dir=temp_config_dir, suffix=fileExtname) - if platform.system() == 'Windows': + temp_config_file = tempfile.NamedTemporaryFile(dir=temp_config_dir, suffix=fileExtname) + if platform.system() == "Windows": temp_config_file.close() temp_config_name = osp.basename(temp_config_file.name) # Substitute predefined variables if use_predefined_variables: - Config._substitute_predefined_vars(filename, - temp_config_file.name) + Config._substitute_predefined_vars(filename, temp_config_file.name) else: shutil.copyfile(filename, temp_config_file.name) # Substitute base variables from placeholders to strings - base_var_dict = Config._pre_substitute_base_vars( - temp_config_file.name, temp_config_file.name) + base_var_dict = Config._pre_substitute_base_vars(temp_config_file.name, temp_config_file.name) - if filename.endswith('.py'): + if filename.endswith(".py"): temp_module_name = osp.splitext(temp_config_name)[0] sys.path.insert(0, temp_config_dir) Config._validate_py_syntax(filename) mod = import_module(temp_module_name) sys.path.pop(0) - cfg_dict = { - name: value - for name, value in mod.__dict__.items() - if not name.startswith('__') - } + cfg_dict = {name: value for name, value in mod.__dict__.items() if not name.startswith("__")} # delete imported module del sys.modules[temp_module_name] - elif filename.endswith(('.yml', '.yaml', '.json')): - import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv + elif filename.endswith((".yml", ".yaml", ".json")): + import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv + cfg_dict = mmcv.load(temp_config_file.name) # close temp file temp_config_file.close() @@ -220,26 +206,22 @@ def _file2dict(filename, use_predefined_variables=True): # check deprecation information if DEPRECATION_KEY in cfg_dict: deprecation_info = cfg_dict.pop(DEPRECATION_KEY) - warning_msg = f'The config file {filename} will be deprecated ' \ - 'in the future.' - if 'expected' in deprecation_info: - warning_msg += f' Please use {deprecation_info["expected"]} ' \ - 'instead.' - if 'reference' in deprecation_info: - warning_msg += ' More information can be found at ' \ - f'{deprecation_info["reference"]}' + warning_msg = f"The config file {filename} will be deprecated " "in the future." + if "expected" in deprecation_info: + warning_msg += f' Please use {deprecation_info["expected"]} ' "instead." + if "reference" in deprecation_info: + warning_msg += " More information can be found at " f'{deprecation_info["reference"]}' warnings.warn(warning_msg) - cfg_text = filename + '\n' - with open(filename, 'r', encoding='utf-8') as f: + cfg_text = filename + "\n" + with open(filename, "r", encoding="utf-8") as f: # Setting encoding explicitly to resolve coding issue on windows cfg_text += f.read() if BASE_KEY in cfg_dict: cfg_dir = osp.dirname(filename) base_filename = cfg_dict.pop(BASE_KEY) - base_filename = base_filename if isinstance( - base_filename, list) else [base_filename] + base_filename = base_filename if isinstance(base_filename, list) else [base_filename] cfg_dict_list = list() cfg_text_list = list() @@ -252,20 +234,18 @@ def _file2dict(filename, use_predefined_variables=True): for c in cfg_dict_list: duplicate_keys = base_cfg_dict.keys() & c.keys() if len(duplicate_keys) > 0: - raise KeyError('Duplicate key is not allowed among bases. ' - f'Duplicate keys: {duplicate_keys}') + raise KeyError("Duplicate key is not allowed among bases. " f"Duplicate keys: {duplicate_keys}") base_cfg_dict.update(c) # Substitute base variables from strings to their actual values - cfg_dict = Config._substitute_base_vars(cfg_dict, base_var_dict, - base_cfg_dict) + cfg_dict = Config._substitute_base_vars(cfg_dict, base_var_dict, base_cfg_dict) base_cfg_dict = Config._merge_a_into_b(cfg_dict, base_cfg_dict) cfg_dict = base_cfg_dict # merge cfg_text cfg_text_list.append(cfg_text) - cfg_text = '\n'.join(cfg_text_list) + cfg_text = "\n".join(cfg_text_list) return cfg_dict, cfg_text @@ -307,30 +287,27 @@ def _merge_a_into_b(a, b, allow_list_keys=False): if allow_list_keys and k.isdigit() and isinstance(b, list): k = int(k) if len(b) <= k: - raise KeyError(f'Index {k} exceeds the length of list {b}') + raise KeyError(f"Index {k} exceeds the length of list {b}") b[k] = Config._merge_a_into_b(v, b[k], allow_list_keys) - elif isinstance(v, - dict) and k in b and not v.pop(DELETE_KEY, False): + elif isinstance(v, dict) and k in b and not v.pop(DELETE_KEY, False): allowed_types = (dict, list) if allow_list_keys else dict if not isinstance(b[k], allowed_types): raise TypeError( - f'{k}={v} in child config cannot inherit from base ' - f'because {k} is a dict in the child config but is of ' - f'type {type(b[k])} in base config. You may set ' - f'`{DELETE_KEY}=True` to ignore the base config') + f"{k}={v} in child config cannot inherit from base " + f"because {k} is a dict in the child config but is of " + f"type {type(b[k])} in base config. You may set " + f"`{DELETE_KEY}=True` to ignore the base config" + ) b[k] = Config._merge_a_into_b(v, b[k], allow_list_keys) else: b[k] = v return b @staticmethod - def fromfile(filename, - use_predefined_variables=True, - import_custom_modules=True): - cfg_dict, cfg_text = Config._file2dict(filename, - use_predefined_variables) - if import_custom_modules and cfg_dict.get('custom_imports', None): - import_modules_from_strings(**cfg_dict['custom_imports']) + def fromfile(filename, use_predefined_variables=True, import_custom_modules=True): + cfg_dict, cfg_text = Config._file2dict(filename, use_predefined_variables) + if import_custom_modules and cfg_dict.get("custom_imports", None): + import_modules_from_strings(**cfg_dict["custom_imports"]) return Config(cfg_dict, cfg_text=cfg_text, filename=filename) @staticmethod @@ -345,15 +322,12 @@ def fromstring(cfg_str, file_format): Returns: obj:`Config`: Config obj. """ - if file_format not in ['.py', '.json', '.yaml', '.yml']: - raise IOError('Only py/yml/yaml/json type are supported now!') - if file_format != '.py' and 'dict(' in cfg_str: + if file_format not in [".py", ".json", ".yaml", ".yml"]: + raise IOError("Only py/yml/yaml/json type are supported now!") + if file_format != ".py" and "dict(" in cfg_str: # check if users specify a wrong suffix for python - warnings.warn( - 'Please check "file_format", the file format may be .py') - with tempfile.NamedTemporaryFile( - 'w', encoding='utf-8', suffix=file_format, - delete=False) as temp_file: + warnings.warn('Please check "file_format", the file format may be .py') + with tempfile.NamedTemporaryFile("w", encoding="utf-8", suffix=file_format, delete=False) as temp_file: temp_file.write(cfg_str) # on windows, previous implementation cause error # see PR 1077 for details @@ -365,11 +339,11 @@ def fromstring(cfg_str, file_format): def auto_argparser(description=None): """Generate argparser from config file automatically (experimental)""" partial_parser = ArgumentParser(description=description) - partial_parser.add_argument('config', help='config file path') + partial_parser.add_argument("config", help="config file path") cfg_file = partial_parser.parse_known_args()[0].config cfg = Config.fromfile(cfg_file) parser = ArgumentParser(description=description) - parser.add_argument('config', help='config file path') + parser.add_argument("config", help="config file path") add_args(parser, cfg) return parser, cfg @@ -377,22 +351,21 @@ def __init__(self, cfg_dict=None, cfg_text=None, filename=None): if cfg_dict is None: cfg_dict = dict() elif not isinstance(cfg_dict, dict): - raise TypeError('cfg_dict must be a dict, but ' - f'got {type(cfg_dict)}') + raise TypeError("cfg_dict must be a dict, but " f"got {type(cfg_dict)}") for key in cfg_dict: if key in RESERVED_KEYS: - raise KeyError(f'{key} is reserved for config file') + raise KeyError(f"{key} is reserved for config file") - super(Config, self).__setattr__('_cfg_dict', ConfigDict(cfg_dict)) - super(Config, self).__setattr__('_filename', filename) + super(Config, self).__setattr__("_cfg_dict", ConfigDict(cfg_dict)) + super(Config, self).__setattr__("_filename", filename) if cfg_text: text = cfg_text elif filename: - with open(filename, 'r') as f: + with open(filename, "r") as f: text = f.read() else: - text = '' - super(Config, self).__setattr__('_text', text) + text = "" + super(Config, self).__setattr__("_text", text) @property def filename(self): @@ -404,17 +377,16 @@ def text(self): @property def pretty_text(self): - indent = 4 def _indent(s_, num_spaces): - s = s_.split('\n') + s = s_.split("\n") if len(s) == 1: return s_ first = s.pop(0) - s = [(num_spaces * ' ') + line for line in s] - s = '\n'.join(s) - s = first + '\n' + s + s = [(num_spaces * " ") + line for line in s] + s = "\n".join(s) + s = first + "\n" + s return s def _format_basic_types(k, v, use_mapping=False): @@ -425,9 +397,9 @@ def _format_basic_types(k, v, use_mapping=False): if use_mapping: k_str = f"'{k}'" if isinstance(k, str) else str(k) - attr_str = f'{k_str}: {v_str}' + attr_str = f"{k_str}: {v_str}" else: - attr_str = f'{str(k)}={v_str}' + attr_str = f"{str(k)}={v_str}" attr_str = _indent(attr_str, indent) return attr_str @@ -435,16 +407,14 @@ def _format_basic_types(k, v, use_mapping=False): def _format_list(k, v, use_mapping=False): # check if all items in the list are dict if all(isinstance(_, dict) for _ in v): - v_str = '[\n' - v_str += '\n'.join( - f'dict({_indent(_format_dict(v_), indent)}),' - for v_ in v).rstrip(',') + v_str = "[\n" + v_str += "\n".join(f"dict({_indent(_format_dict(v_), indent)})," for v_ in v).rstrip(",") if use_mapping: k_str = f"'{k}'" if isinstance(k, str) else str(k) - attr_str = f'{k_str}: {v_str}' + attr_str = f"{k_str}: {v_str}" else: - attr_str = f'{str(k)}={v_str}' - attr_str = _indent(attr_str, indent) + ']' + attr_str = f"{str(k)}={v_str}" + attr_str = _indent(attr_str, indent) + "]" else: attr_str = _format_basic_types(k, v, use_mapping) return attr_str @@ -452,37 +422,36 @@ def _format_list(k, v, use_mapping=False): def _contain_invalid_identifier(dict_str): contain_invalid_identifier = False for key_name in dict_str: - contain_invalid_identifier |= \ - (not str(key_name).isidentifier()) + contain_invalid_identifier |= not str(key_name).isidentifier() return contain_invalid_identifier def _format_dict(input_dict, outest_level=False): - r = '' + r = "" s = [] use_mapping = _contain_invalid_identifier(input_dict) if use_mapping: - r += '{' + r += "{" for idx, (k, v) in enumerate(input_dict.items()): is_last = idx >= len(input_dict) - 1 - end = '' if outest_level or is_last else ',' + end = "" if outest_level or is_last else "," if isinstance(v, dict): - v_str = '\n' + _format_dict(v) + v_str = "\n" + _format_dict(v) if use_mapping: k_str = f"'{k}'" if isinstance(k, str) else str(k) - attr_str = f'{k_str}: dict({v_str}' + attr_str = f"{k_str}: dict({v_str}" else: - attr_str = f'{str(k)}=dict({v_str}' - attr_str = _indent(attr_str, indent) + ')' + end + attr_str = f"{str(k)}=dict({v_str}" + attr_str = _indent(attr_str, indent) + ")" + end elif isinstance(v, list): attr_str = _format_list(k, v, use_mapping) + end else: attr_str = _format_basic_types(k, v, use_mapping) + end s.append(attr_str) - r += '\n'.join(s) + r += "\n".join(s) if use_mapping: - r += '}' + r += "}" return r cfg_dict = self._cfg_dict.to_dict() @@ -491,7 +460,7 @@ def _format_dict(input_dict, outest_level=False): return text def __repr__(self): - return f'Config (path: {self.filename}): {self._cfg_dict.__repr__()}' + return f"Config (path: {self.filename}): {self._cfg_dict.__repr__()}" def __len__(self): return len(self._cfg_dict) @@ -520,22 +489,23 @@ def __getstate__(self): def __setstate__(self, state): _cfg_dict, _filename, _text = state - super(Config, self).__setattr__('_cfg_dict', _cfg_dict) - super(Config, self).__setattr__('_filename', _filename) - super(Config, self).__setattr__('_text', _text) + super(Config, self).__setattr__("_cfg_dict", _cfg_dict) + super(Config, self).__setattr__("_filename", _filename) + super(Config, self).__setattr__("_text", _text) def dump(self, file=None): - cfg_dict = super(Config, self).__getattribute__('_cfg_dict').to_dict() - if self.filename.endswith('.py'): + cfg_dict = super(Config, self).__getattribute__("_cfg_dict").to_dict() + if self.filename.endswith(".py"): if file is None: return self.pretty_text else: - with open(file, 'w', encoding='utf-8') as f: + with open(file, "w", encoding="utf-8") as f: f.write(self.pretty_text) else: - import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv + import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv + if file is None: - file_format = self.filename.split('.')[-1] + file_format = self.filename.split(".")[-1] return mmcv.dump(cfg_dict, file_format=file_format) else: mmcv.dump(cfg_dict, file) @@ -573,18 +543,17 @@ def merge_from_dict(self, options, allow_list_keys=True): option_cfg_dict = {} for full_key, v in options.items(): d = option_cfg_dict - key_list = full_key.split('.') + key_list = full_key.split(".") for subkey in key_list[:-1]: d.setdefault(subkey, ConfigDict()) d = d[subkey] subkey = key_list[-1] d[subkey] = v - cfg_dict = super(Config, self).__getattribute__('_cfg_dict') + cfg_dict = super(Config, self).__getattribute__("_cfg_dict") super(Config, self).__setattr__( - '_cfg_dict', - Config._merge_a_into_b( - option_cfg_dict, cfg_dict, allow_list_keys=allow_list_keys)) + "_cfg_dict", Config._merge_a_into_b(option_cfg_dict, cfg_dict, allow_list_keys=allow_list_keys) + ) class DictAction(Action): @@ -606,8 +575,8 @@ def _parse_int_float_bool(val): return float(val) except ValueError: pass - if val.lower() in ['true', 'false']: - return True if val.lower() == 'true' else False + if val.lower() in ["true", "false"]: + return True if val.lower() == "true" else False return val @staticmethod @@ -638,28 +607,27 @@ def find_next_comma(string): chars inside '()' and '[]' are treated as one element and thus ',' inside these brackets are ignored. """ - assert (string.count('(') == string.count(')')) and ( - string.count('[') == string.count(']')), \ - f'Imbalanced brackets exist in {string}' + assert (string.count("(") == string.count(")")) and ( + string.count("[") == string.count("]") + ), f"Imbalanced brackets exist in {string}" end = len(string) for idx, char in enumerate(string): pre = string[:idx] # The string before this ',' is balanced - if ((char == ',') and (pre.count('(') == pre.count(')')) - and (pre.count('[') == pre.count(']'))): + if (char == ",") and (pre.count("(") == pre.count(")")) and (pre.count("[") == pre.count("]")): end = idx break return end # Strip ' and " characters and replace whitespace. - val = val.strip('\'\"').replace(' ', '') + val = val.strip("'\"").replace(" ", "") is_tuple = False - if val.startswith('(') and val.endswith(')'): + if val.startswith("(") and val.endswith(")"): is_tuple = True val = val[1:-1] - elif val.startswith('[') and val.endswith(']'): + elif val.startswith("[") and val.endswith("]"): val = val[1:-1] - elif ',' not in val: + elif "," not in val: # val is a single value return DictAction._parse_int_float_bool(val) @@ -668,7 +636,7 @@ def find_next_comma(string): comma_idx = find_next_comma(val) element = DictAction._parse_iterable(val[:comma_idx]) values.append(element) - val = val[comma_idx + 1:] + val = val[comma_idx + 1 :] if is_tuple: values = tuple(values) return values @@ -676,6 +644,6 @@ def find_next_comma(string): def __call__(self, parser, namespace, values, option_string=None): options = {} for kv in values: - key, val = kv.split('=', maxsplit=1) + key, val = kv.split("=", maxsplit=1) options[key] = self._parse_iterable(val) setattr(namespace, self.dest, options) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/env.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/env.py index a0418965..ac785392 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/env.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/env.py @@ -9,7 +9,7 @@ import cv2 import torch -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv from .parrots_wrapper import get_build_config @@ -36,60 +36,64 @@ def collect_env(): - MMCV CUDA Compiler: The CUDA version for compiling MMCV ops. """ env_info = {} - env_info['sys.platform'] = sys.platform - env_info['Python'] = sys.version.replace('\n', '') + env_info["sys.platform"] = sys.platform + env_info["Python"] = sys.version.replace("\n", "") cuda_available = torch.cuda.is_available() - env_info['CUDA available'] = cuda_available + env_info["CUDA available"] = cuda_available if cuda_available: devices = defaultdict(list) for k in range(torch.cuda.device_count()): devices[torch.cuda.get_device_name(k)].append(str(k)) for name, device_ids in devices.items(): - env_info['GPU ' + ','.join(device_ids)] = name + env_info["GPU " + ",".join(device_ids)] = name + + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _get_cuda_home - from comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _get_cuda_home CUDA_HOME = _get_cuda_home() - env_info['CUDA_HOME'] = CUDA_HOME + env_info["CUDA_HOME"] = CUDA_HOME if CUDA_HOME is not None and osp.isdir(CUDA_HOME): try: - nvcc = osp.join(CUDA_HOME, 'bin/nvcc') - nvcc = subprocess.check_output( - f'"{nvcc}" -V | tail -n1', shell=True) - nvcc = nvcc.decode('utf-8').strip() + nvcc = osp.join(CUDA_HOME, "bin/nvcc") + nvcc = subprocess.check_output(f'"{nvcc}" -V | tail -n1', shell=True) + nvcc = nvcc.decode("utf-8").strip() except subprocess.SubprocessError: - nvcc = 'Not Available' - env_info['NVCC'] = nvcc + nvcc = "Not Available" + env_info["NVCC"] = nvcc try: - gcc = subprocess.check_output('gcc --version | head -n1', shell=True) - gcc = gcc.decode('utf-8').strip() - env_info['GCC'] = gcc + gcc = subprocess.check_output("gcc --version | head -n1", shell=True) + gcc = gcc.decode("utf-8").strip() + env_info["GCC"] = gcc except subprocess.CalledProcessError: # gcc is unavailable - env_info['GCC'] = 'n/a' + env_info["GCC"] = "n/a" - env_info['PyTorch'] = torch.__version__ - env_info['PyTorch compiling details'] = get_build_config() + env_info["PyTorch"] = torch.__version__ + env_info["PyTorch compiling details"] = get_build_config() try: import torchvision - env_info['TorchVision'] = torchvision.__version__ + + env_info["TorchVision"] = torchvision.__version__ except ModuleNotFoundError: pass - env_info['OpenCV'] = cv2.__version__ + env_info["OpenCV"] = cv2.__version__ - env_info['MMCV'] = mmcv.__version__ + env_info["MMCV"] = mmcv.__version__ try: - from comfy_controlnet_preprocessors.uniformer.mmcv.ops import get_compiler_version, get_compiling_cuda_version + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.ops import ( + get_compiler_version, + get_compiling_cuda_version, + ) except ModuleNotFoundError: - env_info['MMCV Compiler'] = 'n/a' - env_info['MMCV CUDA Compiler'] = 'n/a' + env_info["MMCV Compiler"] = "n/a" + env_info["MMCV CUDA Compiler"] = "n/a" else: - env_info['MMCV Compiler'] = get_compiler_version() - env_info['MMCV CUDA Compiler'] = get_compiling_cuda_version() + env_info["MMCV Compiler"] = get_compiler_version() + env_info["MMCV CUDA Compiler"] = get_compiling_cuda_version() return env_info diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/ext_loader.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/ext_loader.py index 2ba00dc4..30a9720b 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/ext_loader.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/ext_loader.py @@ -7,54 +7,54 @@ import torch -if torch.__version__ != 'parrots': +if torch.__version__ != "parrots": def load_ext(name, funcs): - ext = importlib.import_module('comfy_controlnet_preprocessors.uniformer.mmcv.' + name) + ext = importlib.import_module("hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv." + name) for fun in funcs: - assert hasattr(ext, fun), f'{fun} miss in module {name}' + assert hasattr(ext, fun), f"{fun} miss in module {name}" return ext + else: from parrots import extension from parrots.base import ParrotsException has_return_value_ops = [ - 'nms', - 'softnms', - 'nms_match', - 'nms_rotated', - 'top_pool_forward', - 'top_pool_backward', - 'bottom_pool_forward', - 'bottom_pool_backward', - 'left_pool_forward', - 'left_pool_backward', - 'right_pool_forward', - 'right_pool_backward', - 'fused_bias_leakyrelu', - 'upfirdn2d', - 'ms_deform_attn_forward', - 'pixel_group', - 'contour_expand', + "nms", + "softnms", + "nms_match", + "nms_rotated", + "top_pool_forward", + "top_pool_backward", + "bottom_pool_forward", + "bottom_pool_backward", + "left_pool_forward", + "left_pool_backward", + "right_pool_forward", + "right_pool_backward", + "fused_bias_leakyrelu", + "upfirdn2d", + "ms_deform_attn_forward", + "pixel_group", + "contour_expand", ] def get_fake_func(name, e): - def fake_func(*args, **kwargs): - warnings.warn(f'{name} is not supported in parrots now') + warnings.warn(f"{name} is not supported in parrots now") raise e return fake_func def load_ext(name, funcs): - ExtModule = namedtuple('ExtModule', funcs) + ExtModule = namedtuple("ExtModule", funcs) ext_list = [] lib_root = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) for fun in funcs: try: ext_fun = extension.load(fun, name, lib_dir=lib_root) except ParrotsException as e: - if 'No element registered' not in e.message: + if "No element registered" not in e.message: warnings.warn(e.message) ext_fun = get_fake_func(fun, e) ext_list.append(ext_fun) @@ -67,5 +67,5 @@ def load_ext(name, funcs): def check_ops_exist(): - ext_loader = pkgutil.find_loader('mmcv._ext') + ext_loader = pkgutil.find_loader("mmcv._ext") return ext_loader is not None diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/timer.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/timer.py index 4f6abf4b..c3146835 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/timer.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/timer.py @@ -3,7 +3,6 @@ class TimerError(Exception): - def __init__(self, message): self.message = message super(TimerError, self).__init__(message) @@ -15,7 +14,7 @@ class Timer: :Example: >>> import time - >>> import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv + >>> import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv >>> with mmcv.Timer(): >>> # simulate a code block that will run for 1s >>> time.sleep(1) @@ -37,7 +36,7 @@ class Timer: def __init__(self, start=True, print_tmpl=None): self._is_running = False - self.print_tmpl = print_tmpl if print_tmpl else '{:.3f}' + self.print_tmpl = print_tmpl if print_tmpl else "{:.3f}" if start: self.start() @@ -67,7 +66,7 @@ def since_start(self): Returns (float): Time in seconds. """ if not self._is_running: - raise TimerError('timer is not running') + raise TimerError("timer is not running") self._t_last = time() return self._t_last - self._t_start @@ -80,7 +79,7 @@ def since_last_check(self): Returns (float): Time in seconds. """ if not self._is_running: - raise TimerError('timer is not running') + raise TimerError("timer is not running") dur = time() - self._t_last self._t_last = time() return dur @@ -98,7 +97,7 @@ def check_time(timer_id): :Example: >>> import time - >>> import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv + >>> import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv >>> for i in range(1, 6): >>> # simulate a code block >>> time.sleep(i) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/trace.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/trace.py index 57b1dc99..857e9ad8 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/trace.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/utils/trace.py @@ -2,12 +2,11 @@ import torch -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import digit_version +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import digit_version def is_jit_tracing() -> bool: - if (torch.__version__ != 'parrots' - and digit_version(torch.__version__) >= digit_version('1.6.0')): + if torch.__version__ != "parrots" and digit_version(torch.__version__) >= digit_version("1.6.0"): on_trace = torch.jit.is_tracing() # In PyTorch 1.6, torch.jit.is_tracing has a bug. # Refers to https://github.com/pytorch/pytorch/issues/42448 @@ -17,7 +16,9 @@ def is_jit_tracing() -> bool: return torch._C._is_tracing() else: warnings.warn( - 'torch.jit.is_tracing is only supported after v1.6.0. ' - 'Therefore is_tracing returns False automatically. Please ' - 'set on_trace manually if you are using trace.', UserWarning) + "torch.jit.is_tracing is only supported after v1.6.0. " + "Therefore is_tracing returns False automatically. Please " + "set on_trace manually if you are using trace.", + UserWarning, + ) return False diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/io.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/io.py index e5ed1ce4..90c8ad08 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/io.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/io.py @@ -3,21 +3,30 @@ from collections import OrderedDict import cv2 -from cv2 import (CAP_PROP_FOURCC, CAP_PROP_FPS, CAP_PROP_FRAME_COUNT, - CAP_PROP_FRAME_HEIGHT, CAP_PROP_FRAME_WIDTH, - CAP_PROP_POS_FRAMES, VideoWriter_fourcc) - -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import (check_file_exist, mkdir_or_exist, scandir, - track_progress) +from cv2 import ( + CAP_PROP_FOURCC, + CAP_PROP_FPS, + CAP_PROP_FRAME_COUNT, + CAP_PROP_FRAME_HEIGHT, + CAP_PROP_FRAME_WIDTH, + CAP_PROP_POS_FRAMES, + VideoWriter_fourcc, +) + +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import ( + check_file_exist, + mkdir_or_exist, + scandir, + track_progress, +) class Cache: - def __init__(self, capacity): self._cache = OrderedDict() self._capacity = int(capacity) if capacity <= 0: - raise ValueError('capacity must be a positive integer') + raise ValueError("capacity must be a positive integer") @property def capacity(self): @@ -52,7 +61,7 @@ class VideoReader: :Example: - >>> import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv + >>> import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv >>> v = mmcv.VideoReader('sample.mp4') >>> len(v) # get the total frame number with `len()` 120 @@ -63,8 +72,8 @@ class VideoReader: def __init__(self, filename, cache_capacity=10): # Check whether the video path is a url - if not filename.startswith(('https://', 'http://')): - check_file_exist(filename, 'Video file not found: ' + filename) + if not filename.startswith(("https://", "http://")): + check_file_exist(filename, "Video file not found: " + filename) self._vcap = cv2.VideoCapture(filename) assert cache_capacity > 0 self._cache = Cache(cache_capacity) @@ -167,8 +176,7 @@ def get_frame(self, frame_id): ndarray or None: Return the frame if successful, otherwise None. """ if frame_id < 0 or frame_id >= self._frame_cnt: - raise IndexError( - f'"frame_id" must be between 0 and {self._frame_cnt - 1}') + raise IndexError(f'"frame_id" must be between 0 and {self._frame_cnt - 1}') if frame_id == self._position: return self.read() if self._cache: @@ -195,13 +203,7 @@ def current_frame(self): return None return self._cache.get(self._position - 1) - def cvt2frames(self, - frame_dir, - file_start=0, - filename_tmpl='{:06d}.jpg', - start=0, - max_num=0, - show_progress=True): + def cvt2frames(self, frame_dir, file_start=0, filename_tmpl="{:06d}.jpg", start=0, max_num=0, show_progress=True): """Convert a video to frame images. Args: @@ -219,7 +221,7 @@ def cvt2frames(self, else: task_num = min(self.frame_cnt - start, max_num) if task_num <= 0: - raise ValueError('start must be less than total frame number') + raise ValueError("start must be less than total frame number") if start > 0: self._set_real_position(start) @@ -231,8 +233,7 @@ def write_frame(file_idx): cv2.imwrite(filename, img) if show_progress: - track_progress(write_frame, range(file_start, - file_start + task_num)) + track_progress(write_frame, range(file_start, file_start + task_num)) else: for i in range(task_num): write_frame(file_start + i) @@ -242,15 +243,12 @@ def __len__(self): def __getitem__(self, index): if isinstance(index, slice): - return [ - self.get_frame(i) - for i in range(*index.indices(self.frame_cnt)) - ] + return [self.get_frame(i) for i in range(*index.indices(self.frame_cnt))] # support negative indexing if index < 0: index += self.frame_cnt if index < 0: - raise IndexError('index out of range') + raise IndexError("index out of range") return self.get_frame(index) def __iter__(self): @@ -273,14 +271,9 @@ def __exit__(self, exc_type, exc_value, traceback): self._vcap.release() -def frames2video(frame_dir, - video_file, - fps=30, - fourcc='XVID', - filename_tmpl='{:06d}.jpg', - start=0, - end=0, - show_progress=True): +def frames2video( + frame_dir, video_file, fps=30, fourcc="XVID", filename_tmpl="{:06d}.jpg", start=0, end=0, show_progress=True +): """Read the frame images from a directory and join them as a video. Args: @@ -295,15 +288,14 @@ def frames2video(frame_dir, show_progress (bool): Whether to show a progress bar. """ if end == 0: - ext = filename_tmpl.split('.')[-1] + ext = filename_tmpl.split(".")[-1] end = len([name for name in scandir(frame_dir, ext)]) first_file = osp.join(frame_dir, filename_tmpl.format(start)) - check_file_exist(first_file, 'The start frame not found: ' + first_file) + check_file_exist(first_file, "The start frame not found: " + first_file) img = cv2.imread(first_file) height, width = img.shape[:2] resolution = (width, height) - vwriter = cv2.VideoWriter(video_file, VideoWriter_fourcc(*fourcc), fps, - resolution) + vwriter = cv2.VideoWriter(video_file, VideoWriter_fourcc(*fourcc), fps, resolution) def write_frame(file_idx): filename = osp.join(frame_dir, filename_tmpl.format(file_idx)) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/optflow.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/optflow.py index 0b2e106a..2cc8a750 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/optflow.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/optflow.py @@ -4,9 +4,9 @@ import cv2 import numpy as np -from comfy_controlnet_preprocessors.uniformer.mmcv.arraymisc import dequantize, quantize -from comfy_controlnet_preprocessors.uniformer.mmcv.image import imread, imwrite -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import is_str +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.arraymisc import dequantize, quantize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.image import imread, imwrite +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import is_str def flowread(flow_or_path, quantize=False, concat_axis=0, *args, **kwargs): @@ -24,33 +24,29 @@ def flowread(flow_or_path, quantize=False, concat_axis=0, *args, **kwargs): """ if isinstance(flow_or_path, np.ndarray): if (flow_or_path.ndim != 3) or (flow_or_path.shape[-1] != 2): - raise ValueError(f'Invalid flow with shape {flow_or_path.shape}') + raise ValueError(f"Invalid flow with shape {flow_or_path.shape}") return flow_or_path elif not is_str(flow_or_path): - raise TypeError(f'"flow_or_path" must be a filename or numpy array, ' - f'not {type(flow_or_path)}') + raise TypeError(f'"flow_or_path" must be a filename or numpy array, ' f"not {type(flow_or_path)}") if not quantize: - with open(flow_or_path, 'rb') as f: + with open(flow_or_path, "rb") as f: try: - header = f.read(4).decode('utf-8') + header = f.read(4).decode("utf-8") except Exception: - raise IOError(f'Invalid flow file: {flow_or_path}') + raise IOError(f"Invalid flow file: {flow_or_path}") else: - if header != 'PIEH': - raise IOError(f'Invalid flow file: {flow_or_path}, ' - 'header does not contain PIEH') + if header != "PIEH": + raise IOError(f"Invalid flow file: {flow_or_path}, " "header does not contain PIEH") w = np.fromfile(f, np.int32, 1).squeeze() h = np.fromfile(f, np.int32, 1).squeeze() flow = np.fromfile(f, np.float32, w * h * 2).reshape((h, w, 2)) else: assert concat_axis in [0, 1] - cat_flow = imread(flow_or_path, flag='unchanged') + cat_flow = imread(flow_or_path, flag="unchanged") if cat_flow.ndim != 2: - raise IOError( - f'{flow_or_path} is not a valid quantized flow file, ' - f'its dimension is {cat_flow.ndim}.') + raise IOError(f"{flow_or_path} is not a valid quantized flow file, " f"its dimension is {cat_flow.ndim}.") assert cat_flow.shape[concat_axis] % 2 == 0 dx, dy = np.split(cat_flow, 2, axis=concat_axis) flow = dequantize_flow(dx, dy, *args, **kwargs) @@ -75,8 +71,8 @@ def flowwrite(flow, filename, quantize=False, concat_axis=0, *args, **kwargs): can be either 0 or 1. Ignored if quantize is False. """ if not quantize: - with open(filename, 'wb') as f: - f.write('PIEH'.encode('utf-8')) + with open(filename, "wb") as f: + f.write("PIEH".encode("utf-8")) np.array([flow.shape[1], flow.shape[0]], dtype=np.int32).tofile(f) flow = flow.astype(np.float32) flow.tofile(f) @@ -110,9 +106,7 @@ def quantize_flow(flow, max_val=0.02, norm=True): dx = dx / w # avoid inplace operations dy = dy / h # use 255 levels instead of 256 to make sure 0 is 0 after dequantization. - flow_comps = [ - quantize(d, -max_val, max_val, 255, np.uint8) for d in [dx, dy] - ] + flow_comps = [quantize(d, -max_val, max_val, 255, np.uint8) for d in [dx, dy]] return tuple(flow_comps) @@ -140,7 +134,7 @@ def dequantize_flow(dx, dy, max_val=0.02, denorm=True): return flow -def flow_warp(img, flow, filling_value=0, interpolate_mode='nearest'): +def flow_warp(img, flow, filling_value=0, interpolate_mode="nearest"): """Use flow to warp img. Args: @@ -153,15 +147,13 @@ def flow_warp(img, flow, filling_value=0, interpolate_mode='nearest'): Returns: ndarray: Warped image with the same shape of img """ - warnings.warn('This function is just for prototyping and cannot ' - 'guarantee the computational efficiency.') - assert flow.ndim == 3, 'Flow must be in 3D arrays.' + warnings.warn("This function is just for prototyping and cannot " "guarantee the computational efficiency.") + assert flow.ndim == 3, "Flow must be in 3D arrays." height = flow.shape[0] width = flow.shape[1] channels = img.shape[2] - output = np.ones( - (height, width, channels), dtype=img.dtype) * filling_value + output = np.ones((height, width, channels), dtype=img.dtype) * filling_value grid = np.indices((height, width)).swapaxes(0, 1).swapaxes(1, 2) dx = grid[:, :, 0] + flow[:, :, 1] @@ -170,34 +162,37 @@ def flow_warp(img, flow, filling_value=0, interpolate_mode='nearest'): sy = np.floor(dy).astype(int) valid = (sx >= 0) & (sx < height - 1) & (sy >= 0) & (sy < width - 1) - if interpolate_mode == 'nearest': - output[valid, :] = img[dx[valid].round().astype(int), - dy[valid].round().astype(int), :] - elif interpolate_mode == 'bilinear': + if interpolate_mode == "nearest": + output[valid, :] = img[dx[valid].round().astype(int), dy[valid].round().astype(int), :] + elif interpolate_mode == "bilinear": # dirty walkround for integer positions eps_ = 1e-6 dx, dy = dx + eps_, dy + eps_ - left_top_ = img[np.floor(dx[valid]).astype(int), - np.floor(dy[valid]).astype(int), :] * ( - np.ceil(dx[valid]) - dx[valid])[:, None] * ( - np.ceil(dy[valid]) - dy[valid])[:, None] - left_down_ = img[np.ceil(dx[valid]).astype(int), - np.floor(dy[valid]).astype(int), :] * ( - dx[valid] - np.floor(dx[valid]))[:, None] * ( - np.ceil(dy[valid]) - dy[valid])[:, None] - right_top_ = img[np.floor(dx[valid]).astype(int), - np.ceil(dy[valid]).astype(int), :] * ( - np.ceil(dx[valid]) - dx[valid])[:, None] * ( - dy[valid] - np.floor(dy[valid]))[:, None] - right_down_ = img[np.ceil(dx[valid]).astype(int), - np.ceil(dy[valid]).astype(int), :] * ( - dx[valid] - np.floor(dx[valid]))[:, None] * ( - dy[valid] - np.floor(dy[valid]))[:, None] + left_top_ = ( + img[np.floor(dx[valid]).astype(int), np.floor(dy[valid]).astype(int), :] + * (np.ceil(dx[valid]) - dx[valid])[:, None] + * (np.ceil(dy[valid]) - dy[valid])[:, None] + ) + left_down_ = ( + img[np.ceil(dx[valid]).astype(int), np.floor(dy[valid]).astype(int), :] + * (dx[valid] - np.floor(dx[valid]))[:, None] + * (np.ceil(dy[valid]) - dy[valid])[:, None] + ) + right_top_ = ( + img[np.floor(dx[valid]).astype(int), np.ceil(dy[valid]).astype(int), :] + * (np.ceil(dx[valid]) - dx[valid])[:, None] + * (dy[valid] - np.floor(dy[valid]))[:, None] + ) + right_down_ = ( + img[np.ceil(dx[valid]).astype(int), np.ceil(dy[valid]).astype(int), :] + * (dx[valid] - np.floor(dx[valid]))[:, None] + * (dy[valid] - np.floor(dy[valid]))[:, None] + ) output[valid, :] = left_top_ + left_down_ + right_top_ + right_down_ else: raise NotImplementedError( - 'We only support interpolation modes of nearest and bilinear, ' - f'but got {interpolate_mode}.') + "We only support interpolation modes of nearest and bilinear, " f"but got {interpolate_mode}." + ) return output.astype(img.dtype) @@ -218,15 +213,14 @@ def flow_from_bytes(content): # header in first 4 bytes header = content[:4] - if header.decode('utf-8') != 'PIEH': - raise Exception('Flow file header does not contain PIEH') + if header.decode("utf-8") != "PIEH": + raise Exception("Flow file header does not contain PIEH") # width in second 4 bytes width = np.frombuffer(content[4:], np.int32, 1).squeeze() # height in third 4 bytes height = np.frombuffer(content[8:], np.int32, 1).squeeze() # after first 12 bytes, all bytes are flow - flow = np.frombuffer(content[12:], np.float32, width * height * 2).reshape( - (height, width, 2)) + flow = np.frombuffer(content[12:], np.float32, width * height * 2).reshape((height, width, 2)) return flow diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/processing.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/processing.py index 506efdd5..13f1bf66 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/processing.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/video/processing.py @@ -4,15 +4,11 @@ import subprocess import tempfile -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import requires_executable +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import requires_executable -@requires_executable('ffmpeg') -def convert_video(in_file, - out_file, - print_cmd=False, - pre_options='', - **kwargs): +@requires_executable("ffmpeg") +def convert_video(in_file, out_file, print_cmd=False, pre_options="", **kwargs): """Convert a video with ffmpeg. This provides a general api to ffmpeg, the executed command is:: @@ -35,30 +31,20 @@ def convert_video(in_file, for k, v in kwargs.items(): if isinstance(v, bool): if v: - options.append(f'-{k}') - elif k == 'log_level': - assert v in [ - 'quiet', 'panic', 'fatal', 'error', 'warning', 'info', - 'verbose', 'debug', 'trace' - ] - options.append(f'-loglevel {v}') + options.append(f"-{k}") + elif k == "log_level": + assert v in ["quiet", "panic", "fatal", "error", "warning", "info", "verbose", "debug", "trace"] + options.append(f"-loglevel {v}") else: - options.append(f'-{k} {v}') - cmd = f'ffmpeg -y {pre_options} -i {in_file} {" ".join(options)} ' \ - f'{out_file}' + options.append(f"-{k} {v}") + cmd = f'ffmpeg -y {pre_options} -i {in_file} {" ".join(options)} ' f"{out_file}" if print_cmd: print(cmd) subprocess.call(cmd, shell=True) -@requires_executable('ffmpeg') -def resize_video(in_file, - out_file, - size=None, - ratio=None, - keep_ar=False, - log_level='info', - print_cmd=False): +@requires_executable("ffmpeg") +def resize_video(in_file, out_file, size=None, ratio=None, keep_ar=False, log_level="info", print_cmd=False): """Resize a video. Args: @@ -72,32 +58,24 @@ def resize_video(in_file, print_cmd (bool): Whether to print the final ffmpeg command. """ if size is None and ratio is None: - raise ValueError('expected size or ratio must be specified') + raise ValueError("expected size or ratio must be specified") if size is not None and ratio is not None: - raise ValueError('size and ratio cannot be specified at the same time') - options = {'log_level': log_level} + raise ValueError("size and ratio cannot be specified at the same time") + options = {"log_level": log_level} if size: if not keep_ar: - options['vf'] = f'scale={size[0]}:{size[1]}' + options["vf"] = f"scale={size[0]}:{size[1]}" else: - options['vf'] = f'scale=w={size[0]}:h={size[1]}:' \ - 'force_original_aspect_ratio=decrease' + options["vf"] = f"scale=w={size[0]}:h={size[1]}:" "force_original_aspect_ratio=decrease" else: if not isinstance(ratio, tuple): ratio = (ratio, ratio) - options['vf'] = f'scale="trunc(iw*{ratio[0]}):trunc(ih*{ratio[1]})"' + options["vf"] = f'scale="trunc(iw*{ratio[0]}):trunc(ih*{ratio[1]})"' convert_video(in_file, out_file, print_cmd, **options) -@requires_executable('ffmpeg') -def cut_video(in_file, - out_file, - start=None, - end=None, - vcodec=None, - acodec=None, - log_level='info', - print_cmd=False): +@requires_executable("ffmpeg") +def cut_video(in_file, out_file, start=None, end=None, vcodec=None, acodec=None, log_level="info", print_cmd=False): """Cut a clip from a video. Args: @@ -110,27 +88,22 @@ def cut_video(in_file, log_level (str): Logging level of ffmpeg. print_cmd (bool): Whether to print the final ffmpeg command. """ - options = {'log_level': log_level} + options = {"log_level": log_level} if vcodec is None: - options['vcodec'] = 'copy' + options["vcodec"] = "copy" if acodec is None: - options['acodec'] = 'copy' + options["acodec"] = "copy" if start: - options['ss'] = start + options["ss"] = start else: start = 0 if end: - options['t'] = end - start + options["t"] = end - start convert_video(in_file, out_file, print_cmd, **options) -@requires_executable('ffmpeg') -def concat_video(video_list, - out_file, - vcodec=None, - acodec=None, - log_level='info', - print_cmd=False): +@requires_executable("ffmpeg") +def concat_video(video_list, out_file, vcodec=None, acodec=None, log_level="info", print_cmd=False): """Concatenate multiple videos into a single one. Args: @@ -141,20 +114,15 @@ def concat_video(video_list, log_level (str): Logging level of ffmpeg. print_cmd (bool): Whether to print the final ffmpeg command. """ - tmp_filehandler, tmp_filename = tempfile.mkstemp(suffix='.txt', text=True) - with open(tmp_filename, 'w') as f: + tmp_filehandler, tmp_filename = tempfile.mkstemp(suffix=".txt", text=True) + with open(tmp_filename, "w") as f: for filename in video_list: - f.write(f'file {osp.abspath(filename)}\n') - options = {'log_level': log_level} + f.write(f"file {osp.abspath(filename)}\n") + options = {"log_level": log_level} if vcodec is None: - options['vcodec'] = 'copy' + options["vcodec"] = "copy" if acodec is None: - options['acodec'] = 'copy' - convert_video( - tmp_filename, - out_file, - print_cmd, - pre_options='-f concat -safe 0', - **options) + options["acodec"] = "copy" + convert_video(tmp_filename, out_file, print_cmd, pre_options="-f concat -safe 0", **options) os.close(tmp_filehandler) os.remove(tmp_filename) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/color.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/color.py index 5f59c054..ef13de93 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/color.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/color.py @@ -3,7 +3,7 @@ import numpy as np -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import is_str +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import is_str class Color(Enum): @@ -11,6 +11,7 @@ class Color(Enum): Contains red, green, blue, cyan, yellow, magenta, white and black. """ + red = (0, 0, 255) green = (0, 255, 0) blue = (255, 0, 0) @@ -48,4 +49,4 @@ def color_val(color): color = color.astype(np.uint8) return tuple(color) else: - raise TypeError(f'Invalid type for color: {type(color)}') + raise TypeError(f"Invalid type for color: {type(color)}") diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/image.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/image.py index 8d9a10cb..1bf3befd 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/image.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/image.py @@ -2,11 +2,11 @@ import cv2 import numpy as np -from comfy_controlnet_preprocessors.uniformer.mmcv.image import imread, imwrite +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.image import imread, imwrite from .color import color_val -def imshow(img, win_name='', wait_time=0): +def imshow(img, win_name="", wait_time=0): """Show an image. Args: @@ -27,15 +27,9 @@ def imshow(img, win_name='', wait_time=0): ret = cv2.waitKey(wait_time) -def imshow_bboxes(img, - bboxes, - colors='green', - top_k=-1, - thickness=1, - show=True, - win_name='', - wait_time=0, - out_file=None): +def imshow_bboxes( + img, bboxes, colors="green", top_k=-1, thickness=1, show=True, win_name="", wait_time=0, out_file=None +): """Draw bboxes on an image. Args: @@ -71,8 +65,7 @@ def imshow_bboxes(img, for j in range(_top_k): left_top = (_bboxes[j, 0], _bboxes[j, 1]) right_bottom = (_bboxes[j, 2], _bboxes[j, 3]) - cv2.rectangle( - img, left_top, right_bottom, colors[i], thickness=thickness) + cv2.rectangle(img, left_top, right_bottom, colors[i], thickness=thickness) if show: imshow(img, win_name, wait_time) @@ -81,19 +74,21 @@ def imshow_bboxes(img, return img -def imshow_det_bboxes(img, - bboxes, - labels, - class_names=None, - score_thr=0, - bbox_color='green', - text_color='green', - thickness=1, - font_scale=0.5, - show=True, - win_name='', - wait_time=0, - out_file=None): +def imshow_det_bboxes( + img, + bboxes, + labels, + class_names=None, + score_thr=0, + bbox_color="green", + text_color="green", + thickness=1, + font_scale=0.5, + show=True, + win_name="", + wait_time=0, + out_file=None, +): """Draw bboxes and class labels (with scores) on an image. Args: @@ -136,14 +131,11 @@ def imshow_det_bboxes(img, bbox_int = bbox.astype(np.int32) left_top = (bbox_int[0], bbox_int[1]) right_bottom = (bbox_int[2], bbox_int[3]) - cv2.rectangle( - img, left_top, right_bottom, bbox_color, thickness=thickness) - label_text = class_names[ - label] if class_names is not None else f'cls {label}' + cv2.rectangle(img, left_top, right_bottom, bbox_color, thickness=thickness) + label_text = class_names[label] if class_names is not None else f"cls {label}" if len(bbox) > 4: - label_text += f'|{bbox[-1]:.02f}' - cv2.putText(img, label_text, (bbox_int[0], bbox_int[1] - 2), - cv2.FONT_HERSHEY_COMPLEX, font_scale, text_color) + label_text += f"|{bbox[-1]:.02f}" + cv2.putText(img, label_text, (bbox_int[0], bbox_int[1] - 2), cv2.FONT_HERSHEY_COMPLEX, font_scale, text_color) if show: imshow(img, win_name, wait_time) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/optflow.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/optflow.py index 4e717b03..47e91aa9 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/optflow.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv/visualization/optflow.py @@ -3,12 +3,12 @@ import numpy as np -from comfy_controlnet_preprocessors.uniformer.mmcv.image import rgb2bgr -from comfy_controlnet_preprocessors.uniformer.mmcv.video import flowread +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.image import rgb2bgr +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.video import flowread from .image import imshow -def flowshow(flow, win_name='', wait_time=0): +def flowshow(flow, win_name="", wait_time=0): """Show optical flow. Args: @@ -43,9 +43,7 @@ def flow2rgb(flow, color_wheel=None, unknown_thr=1e6): dx = flow[:, :, 0].copy() dy = flow[:, :, 1].copy() - ignore_inds = ( - np.isnan(dx) | np.isnan(dy) | (np.abs(dx) > unknown_thr) | - (np.abs(dy) > unknown_thr)) + ignore_inds = np.isnan(dx) | np.isnan(dy) | (np.abs(dx) > unknown_thr) | (np.abs(dy) > unknown_thr) dx[ignore_inds] = 0 dy[ignore_inds] = 0 @@ -62,8 +60,7 @@ def flow2rgb(flow, color_wheel=None, unknown_thr=1e6): bin_left = np.floor(bin_real).astype(int) bin_right = (bin_left + 1) % num_bins w = (bin_real - bin_left.astype(np.float32))[..., None] - flow_img = (1 - - w) * color_wheel[bin_left, :] + w * color_wheel[bin_right, :] + flow_img = (1 - w) * color_wheel[bin_left, :] + w * color_wheel[bin_right, :] small_ind = rad <= 1 flow_img[small_ind] = 1 - rad[small_ind, None] * (1 - flow_img[small_ind]) flow_img[np.logical_not(small_ind)] *= 0.75 @@ -106,7 +103,7 @@ def make_color_wheel(bins=None): col = 0 for i, color in enumerate([ry, yg, gc, cb, bm, mr]): for j in range(3): - color_wheel[j, col:col + bins[i]] = color[j] + color_wheel[j, col : col + bins[i]] = color[j] col += bins[i] return color_wheel.T diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv_custom/checkpoint.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv_custom/checkpoint.py index d578a342..6181d56c 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv_custom/checkpoint.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmcv_custom/checkpoint.py @@ -15,24 +15,22 @@ from torch.utils import model_zoo from torch.nn import functional as F -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv -from comfy_controlnet_preprocessors.uniformer.mmcv.fileio import FileClient -from comfy_controlnet_preprocessors.uniformer.mmcv.fileio import load as load_file -from comfy_controlnet_preprocessors.uniformer.mmcv.parallel import is_module_wrapper -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import mkdir_or_exist -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import get_dist_info +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.fileio import FileClient +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.fileio import load as load_file +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.parallel import is_module_wrapper +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import mkdir_or_exist +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import get_dist_info -ENV_MMCV_HOME = 'MMCV_HOME' -ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME' -DEFAULT_CACHE_DIR = '~/.cache' +ENV_MMCV_HOME = "MMCV_HOME" +ENV_XDG_CACHE_HOME = "XDG_CACHE_HOME" +DEFAULT_CACHE_DIR = "~/.cache" def _get_mmcv_home(): mmcv_home = os.path.expanduser( - os.getenv( - ENV_MMCV_HOME, - os.path.join( - os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'mmcv'))) + os.getenv(ENV_MMCV_HOME, os.path.join(os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), "mmcv")) + ) mkdir_or_exist(mmcv_home) return mmcv_home @@ -58,46 +56,40 @@ def load_state_dict(module, state_dict, strict=False, logger=None): all_missing_keys = [] err_msg = [] - metadata = getattr(state_dict, '_metadata', None) + metadata = getattr(state_dict, "_metadata", None) state_dict = state_dict.copy() if metadata is not None: state_dict._metadata = metadata # use _load_from_state_dict to enable checkpoint version control - def load(module, prefix=''): + def load(module, prefix=""): # recursively check parallel module in case that the model has a # complicated structure, e.g., nn.Module(nn.Module(DDP)) if is_module_wrapper(module): module = module.module - local_metadata = {} if metadata is None else metadata.get( - prefix[:-1], {}) - module._load_from_state_dict(state_dict, prefix, local_metadata, True, - all_missing_keys, unexpected_keys, - err_msg) + local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) + module._load_from_state_dict( + state_dict, prefix, local_metadata, True, all_missing_keys, unexpected_keys, err_msg + ) for name, child in module._modules.items(): if child is not None: - load(child, prefix + name + '.') + load(child, prefix + name + ".") load(module) load = None # break load->load reference cycle # ignore "num_batches_tracked" of BN layers - missing_keys = [ - key for key in all_missing_keys if 'num_batches_tracked' not in key - ] + missing_keys = [key for key in all_missing_keys if "num_batches_tracked" not in key] if unexpected_keys: - err_msg.append('unexpected key in source ' - f'state_dict: {", ".join(unexpected_keys)}\n') + err_msg.append("unexpected key in source " f'state_dict: {", ".join(unexpected_keys)}\n') if missing_keys: - err_msg.append( - f'missing keys in source state_dict: {", ".join(missing_keys)}\n') + err_msg.append(f'missing keys in source state_dict: {", ".join(missing_keys)}\n') rank, _ = get_dist_info() if len(err_msg) > 0 and rank == 0: - err_msg.insert( - 0, 'The model and loaded state dict do not match exactly\n') - err_msg = '\n'.join(err_msg) + err_msg.insert(0, "The model and loaded state dict do not match exactly\n") + err_msg = "\n".join(err_msg) if strict: raise RuntimeError(err_msg) elif logger is not None: @@ -110,7 +102,7 @@ def load_url_dist(url, model_dir=None): """In distributed setting, this function only download checkpoint at local rank 0.""" rank, world_size = get_dist_info() - rank = int(os.environ.get('LOCAL_RANK', rank)) + rank = int(os.environ.get("LOCAL_RANK", rank)) if rank == 0: checkpoint = model_zoo.load_url(url, model_dir=model_dir) if world_size > 1: @@ -126,10 +118,9 @@ def load_pavimodel_dist(model_path, map_location=None): try: from pavi import modelcloud except ImportError: - raise ImportError( - 'Please install pavi to load checkpoint from modelcloud.') + raise ImportError("Please install pavi to load checkpoint from modelcloud.") rank, world_size = get_dist_info() - rank = int(os.environ.get('LOCAL_RANK', rank)) + rank = int(os.environ.get("LOCAL_RANK", rank)) if rank == 0: model = modelcloud.get(model_path) with TemporaryDirectory() as tmp_dir: @@ -143,8 +134,7 @@ def load_pavimodel_dist(model_path, map_location=None): with TemporaryDirectory() as tmp_dir: downloaded_file = osp.join(tmp_dir, model.name) model.download(downloaded_file) - checkpoint = torch.load( - downloaded_file, map_location=map_location) + checkpoint = torch.load(downloaded_file, map_location=map_location) return checkpoint @@ -152,10 +142,10 @@ def load_fileclient_dist(filename, backend, map_location): """In distributed setting, this function only download checkpoint at local rank 0.""" rank, world_size = get_dist_info() - rank = int(os.environ.get('LOCAL_RANK', rank)) - allowed_backends = ['ceph'] + rank = int(os.environ.get("LOCAL_RANK", rank)) + allowed_backends = ["ceph"] if backend not in allowed_backends: - raise ValueError(f'Load from Backend {backend} is not supported.') + raise ValueError(f"Load from Backend {backend} is not supported.") if rank == 0: fileclient = FileClient(backend=backend) buffer = io.BytesIO(fileclient.get(filename)) @@ -174,19 +164,19 @@ def get_torchvision_models(): for _, name, ispkg in pkgutil.walk_packages(torchvision.models.__path__): if ispkg: continue - _zoo = import_module(f'torchvision.models.{name}') - if hasattr(_zoo, 'model_urls'): - _urls = getattr(_zoo, 'model_urls') + _zoo = import_module(f"torchvision.models.{name}") + if hasattr(_zoo, "model_urls"): + _urls = getattr(_zoo, "model_urls") model_urls.update(_urls) return model_urls def get_external_models(): mmcv_home = _get_mmcv_home() - default_json_path = osp.join(mmcv.__path__[0], 'model_zoo/open_mmlab.json') + default_json_path = osp.join(mmcv.__path__[0], "model_zoo/open_mmlab.json") default_urls = load_file(default_json_path) assert isinstance(default_urls, dict) - external_json_path = osp.join(mmcv_home, 'open_mmlab.json') + external_json_path = osp.join(mmcv_home, "open_mmlab.json") if osp.exists(external_json_path): external_urls = load_file(external_json_path) assert isinstance(external_urls, dict) @@ -196,15 +186,14 @@ def get_external_models(): def get_mmcls_models(): - mmcls_json_path = osp.join(mmcv.__path__[0], 'model_zoo/mmcls.json') + mmcls_json_path = osp.join(mmcv.__path__[0], "model_zoo/mmcls.json") mmcls_urls = load_file(mmcls_json_path) return mmcls_urls def get_deprecated_model_names(): - deprecate_json_path = osp.join(mmcv.__path__[0], - 'model_zoo/deprecated.json') + deprecate_json_path = osp.join(mmcv.__path__[0], "model_zoo/deprecated.json") deprecate_urls = load_file(deprecate_json_path) assert isinstance(deprecate_urls, dict) @@ -212,10 +201,10 @@ def get_deprecated_model_names(): def _process_mmcls_checkpoint(checkpoint): - state_dict = checkpoint['state_dict'] + state_dict = checkpoint["state_dict"] new_state_dict = OrderedDict() for k, v in state_dict.items(): - if k.startswith('backbone.'): + if k.startswith("backbone."): new_state_dict[k[9:]] = v new_checkpoint = dict(state_dict=new_state_dict) @@ -236,58 +225,53 @@ def _load_checkpoint(filename, map_location=None): OrderedDict storing model weights or a dict containing other information, which depends on the checkpoint. """ - if filename.startswith('modelzoo://'): - warnings.warn('The URL scheme of "modelzoo://" is deprecated, please ' - 'use "torchvision://" instead') + if filename.startswith("modelzoo://"): + warnings.warn('The URL scheme of "modelzoo://" is deprecated, please ' 'use "torchvision://" instead') model_urls = get_torchvision_models() model_name = filename[11:] checkpoint = load_url_dist(model_urls[model_name]) - elif filename.startswith('torchvision://'): + elif filename.startswith("torchvision://"): model_urls = get_torchvision_models() model_name = filename[14:] checkpoint = load_url_dist(model_urls[model_name]) - elif filename.startswith('open-mmlab://'): + elif filename.startswith("open-mmlab://"): model_urls = get_external_models() model_name = filename[13:] deprecated_urls = get_deprecated_model_names() if model_name in deprecated_urls: - warnings.warn(f'open-mmlab://{model_name} is deprecated in favor ' - f'of open-mmlab://{deprecated_urls[model_name]}') + warnings.warn( + f"open-mmlab://{model_name} is deprecated in favor " f"of open-mmlab://{deprecated_urls[model_name]}" + ) model_name = deprecated_urls[model_name] model_url = model_urls[model_name] # check if is url - if model_url.startswith(('http://', 'https://')): + if model_url.startswith(("http://", "https://")): checkpoint = load_url_dist(model_url) else: filename = osp.join(_get_mmcv_home(), model_url) if not osp.isfile(filename): - raise IOError(f'{filename} is not a checkpoint file') + raise IOError(f"{filename} is not a checkpoint file") checkpoint = torch.load(filename, map_location=map_location) - elif filename.startswith('mmcls://'): + elif filename.startswith("mmcls://"): model_urls = get_mmcls_models() model_name = filename[8:] checkpoint = load_url_dist(model_urls[model_name]) checkpoint = _process_mmcls_checkpoint(checkpoint) - elif filename.startswith(('http://', 'https://')): + elif filename.startswith(("http://", "https://")): checkpoint = load_url_dist(filename) - elif filename.startswith('pavi://'): + elif filename.startswith("pavi://"): model_path = filename[7:] checkpoint = load_pavimodel_dist(model_path, map_location=map_location) - elif filename.startswith('s3://'): - checkpoint = load_fileclient_dist( - filename, backend='ceph', map_location=map_location) + elif filename.startswith("s3://"): + checkpoint = load_fileclient_dist(filename, backend="ceph", map_location=map_location) else: if not osp.isfile(filename): - raise IOError(f'{filename} is not a checkpoint file') + raise IOError(f"{filename} is not a checkpoint file") checkpoint = torch.load(filename, map_location=map_location) return checkpoint -def load_checkpoint(model, - filename, - map_location='cpu', - strict=False, - logger=None): +def load_checkpoint(model, filename, map_location="cpu", strict=False, logger=None): """Load checkpoint from a file or URI. Args: @@ -306,32 +290,31 @@ def load_checkpoint(model, checkpoint = _load_checkpoint(filename, map_location) # OrderedDict is a subclass of dict if not isinstance(checkpoint, dict): - raise RuntimeError( - f'No state_dict found in checkpoint file {filename}') + raise RuntimeError(f"No state_dict found in checkpoint file {filename}") # get state_dict from checkpoint - if 'state_dict' in checkpoint: - state_dict = checkpoint['state_dict'] - elif 'model' in checkpoint: - state_dict = checkpoint['model'] + if "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] + elif "model" in checkpoint: + state_dict = checkpoint["model"] else: state_dict = checkpoint # strip prefix of state_dict - if list(state_dict.keys())[0].startswith('module.'): + if list(state_dict.keys())[0].startswith("module."): state_dict = {k[7:]: v for k, v in state_dict.items()} # for MoBY, load model of online branch - if sorted(list(state_dict.keys()))[0].startswith('encoder'): - state_dict = {k.replace('encoder.', ''): v for k, v in state_dict.items() if k.startswith('encoder.')} + if sorted(list(state_dict.keys()))[0].startswith("encoder"): + state_dict = {k.replace("encoder.", ""): v for k, v in state_dict.items() if k.startswith("encoder.")} # reshape absolute position embedding - if state_dict.get('absolute_pos_embed') is not None: - absolute_pos_embed = state_dict['absolute_pos_embed'] + if state_dict.get("absolute_pos_embed") is not None: + absolute_pos_embed = state_dict["absolute_pos_embed"] N1, L, C1 = absolute_pos_embed.size() N2, C2, H, W = model.absolute_pos_embed.size() - if N1 != N2 or C1 != C2 or L != H*W: + if N1 != N2 or C1 != C2 or L != H * W: logger.warning("Error in loading absolute_pos_embed, pass") else: - state_dict['absolute_pos_embed'] = absolute_pos_embed.view(N2, H, W, C2).permute(0, 3, 1, 2) + state_dict["absolute_pos_embed"] = absolute_pos_embed.view(N2, H, W, C2).permute(0, 3, 1, 2) # interpolate position bias table if needed relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k] @@ -344,11 +327,11 @@ def load_checkpoint(model, logger.warning(f"Error in loading {table_key}, pass") else: if L1 != L2: - S1 = int(L1 ** 0.5) - S2 = int(L2 ** 0.5) + S1 = int(L1**0.5) + S2 = int(L2**0.5) table_pretrained_resized = F.interpolate( - table_pretrained.permute(1, 0).view(1, nH1, S1, S1), - size=(S2, S2), mode='bicubic') + table_pretrained.permute(1, 0).view(1, nH1, S1, S1), size=(S2, S2), mode="bicubic" + ) state_dict[table_key] = table_pretrained_resized.view(nH2, L2).permute(1, 0) # load state_dict @@ -391,7 +374,7 @@ def _save_to_state_dict(module, destination, prefix, keep_vars): destination[prefix + name] = buf if keep_vars else buf.detach() -def get_state_dict(module, destination=None, prefix='', keep_vars=False): +def get_state_dict(module, destination=None, prefix="", keep_vars=False): """Returns a dictionary containing a whole state of the module. Both parameters and persistent buffers (e.g. running averages) are @@ -421,13 +404,11 @@ def get_state_dict(module, destination=None, prefix='', keep_vars=False): if destination is None: destination = OrderedDict() destination._metadata = OrderedDict() - destination._metadata[prefix[:-1]] = local_metadata = dict( - version=module._version) + destination._metadata[prefix[:-1]] = local_metadata = dict(version=module._version) _save_to_state_dict(module, destination, prefix, keep_vars) for name, child in module._modules.items(): if child is not None: - get_state_dict( - child, destination, prefix + name + '.', keep_vars=keep_vars) + get_state_dict(child, destination, prefix + name + ".", keep_vars=keep_vars) for hook in module._state_dict_hooks.values(): hook_result = hook(module, destination, prefix, local_metadata) if hook_result is not None: @@ -450,35 +431,31 @@ def save_checkpoint(model, filename, optimizer=None, meta=None): if meta is None: meta = {} elif not isinstance(meta, dict): - raise TypeError(f'meta must be a dict or None, but got {type(meta)}') + raise TypeError(f"meta must be a dict or None, but got {type(meta)}") meta.update(mmcv_version=mmcv.__version__, time=time.asctime()) if is_module_wrapper(model): model = model.module - if hasattr(model, 'CLASSES') and model.CLASSES is not None: + if hasattr(model, "CLASSES") and model.CLASSES is not None: # save class name to the meta meta.update(CLASSES=model.CLASSES) - checkpoint = { - 'meta': meta, - 'state_dict': weights_to_cpu(get_state_dict(model)) - } + checkpoint = {"meta": meta, "state_dict": weights_to_cpu(get_state_dict(model))} # save optimizer state dict in the checkpoint if isinstance(optimizer, Optimizer): - checkpoint['optimizer'] = optimizer.state_dict() + checkpoint["optimizer"] = optimizer.state_dict() elif isinstance(optimizer, dict): - checkpoint['optimizer'] = {} + checkpoint["optimizer"] = {} for name, optim in optimizer.items(): - checkpoint['optimizer'][name] = optim.state_dict() + checkpoint["optimizer"][name] = optim.state_dict() - if filename.startswith('pavi://'): + if filename.startswith("pavi://"): try: from pavi import modelcloud from pavi.exception import NodeNotFoundError except ImportError: - raise ImportError( - 'Please install pavi to load checkpoint from modelcloud.') + raise ImportError("Please install pavi to load checkpoint from modelcloud.") model_path = filename[7:] root = modelcloud.Folder() model_dir, model_name = osp.split(model_path) @@ -488,13 +465,13 @@ def save_checkpoint(model, filename, optimizer=None, meta=None): model = root.create_training_model(model_dir) with TemporaryDirectory() as tmp_dir: checkpoint_file = osp.join(tmp_dir, model_name) - with open(checkpoint_file, 'wb') as f: + with open(checkpoint_file, "wb") as f: torch.save(checkpoint, f) f.flush() model.create_file(checkpoint_file, name=model_name) else: mmcv.mkdir_or_exist(osp.dirname(filename)) # immediately flush buffer - with open(filename, 'wb') as f: + with open(filename, "wb") as f: torch.save(checkpoint, f) f.flush() diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/inference.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/inference.py index f0824b07..fb92fcb0 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/inference.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/inference.py @@ -1,13 +1,13 @@ -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import torch -from comfy_controlnet_preprocessors.uniformer.mmcv.parallel import collate, scatter -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.parallel import collate, scatter +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint -from comfy_controlnet_preprocessors.uniformer.mmseg.datasets.pipelines import Compose -from comfy_controlnet_preprocessors.uniformer.mmseg.models import build_segmentor +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.datasets.pipelines import Compose +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.models import build_segmentor -def init_segmentor(config, checkpoint=None, device='cuda:0'): +def init_segmentor(config, checkpoint=None, device="cuda:0"): """Initialize a segmentor from config file. Args: @@ -23,15 +23,14 @@ def init_segmentor(config, checkpoint=None, device='cuda:0'): if isinstance(config, str): config = mmcv.Config.fromfile(config) elif not isinstance(config, mmcv.Config): - raise TypeError('config must be a filename or Config object, ' - 'but got {}'.format(type(config))) + raise TypeError("config must be a filename or Config object, " "but got {}".format(type(config))) config.model.pretrained = None config.model.train_cfg = None - model = build_segmentor(config.model, test_cfg=config.get('test_cfg')) + model = build_segmentor(config.model, test_cfg=config.get("test_cfg")) if checkpoint is not None: - checkpoint = load_checkpoint(model, checkpoint, map_location='cpu') - model.CLASSES = checkpoint['meta']['CLASSES'] - model.PALETTE = checkpoint['meta']['PALETTE'] + checkpoint = load_checkpoint(model, checkpoint, map_location="cpu") + model.CLASSES = checkpoint["meta"]["CLASSES"] + model.PALETTE = checkpoint["meta"]["PALETTE"] model.cfg = config # save the config in the model for convenience model.to(device) model.eval() @@ -52,16 +51,16 @@ def __call__(self, results): dict: ``results`` will be returned containing loaded image. """ - if isinstance(results['img'], str): - results['filename'] = results['img'] - results['ori_filename'] = results['img'] + if isinstance(results["img"], str): + results["filename"] = results["img"] + results["ori_filename"] = results["img"] else: - results['filename'] = None - results['ori_filename'] = None - img = mmcv.imread(results['img']) - results['img'] = img - results['img_shape'] = img.shape - results['ori_shape'] = img.shape + results["filename"] = None + results["ori_filename"] = None + img = mmcv.imread(results["img"]) + results["img"] = img + results["img_shape"] = img.shape + results["ori_shape"] = img.shape return results @@ -89,7 +88,7 @@ def inference_segmentor(model, img): # scatter to specified GPU data = scatter(data, [device])[0] else: - data['img_metas'] = [i.data[0] for i in data['img_metas']] + data["img_metas"] = [i.data[0] for i in data["img_metas"]] # forward the model with torch.no_grad(): @@ -97,14 +96,7 @@ def inference_segmentor(model, img): return result -def show_result_pyplot(model, - img, - result, - palette=None, - fig_size=(15, 10), - opacity=0.5, - title='', - block=True): +def show_result_pyplot(model, img, result, palette=None, fig_size=(15, 10), opacity=0.5, title="", block=True): """Visualize the segmentation results on the image. Args: @@ -123,10 +115,9 @@ def show_result_pyplot(model, block (bool): Whether to block the pyplot figure. Default is True. """ - if hasattr(model, 'module'): + if hasattr(model, "module"): model = model.module - img = model.show_result( - img, result, palette=palette, show=False, opacity=opacity) + img = model.show_result(img, result, palette=palette, show=False, opacity=opacity) # plt.figure(figsize=fig_size) # plt.imshow(mmcv.bgr2rgb(img)) # plt.title(title) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/test.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/test.py index 6d8ebbf1..88a6652a 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/test.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/test.py @@ -3,12 +3,12 @@ import shutil import tempfile -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import numpy as np import torch import torch.distributed as dist -from comfy_controlnet_preprocessors.uniformer.mmcv.image import tensor2imgs -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import get_dist_info +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.image import tensor2imgs +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import get_dist_info def np2tmp(array, temp_file_name=None): @@ -25,18 +25,12 @@ def np2tmp(array, temp_file_name=None): """ if temp_file_name is None: - temp_file_name = tempfile.NamedTemporaryFile( - suffix='.npy', delete=False).name + temp_file_name = tempfile.NamedTemporaryFile(suffix=".npy", delete=False).name np.save(temp_file_name, array) return temp_file_name -def single_gpu_test(model, - data_loader, - show=False, - out_dir=None, - efficient_test=False, - opacity=0.5): +def single_gpu_test(model, data_loader, show=False, out_dir=None, efficient_test=False, opacity=0.5): """Test with single GPU. Args: @@ -63,30 +57,26 @@ def single_gpu_test(model, result = model(return_loss=False, **data) if show or out_dir: - img_tensor = data['img'][0] - img_metas = data['img_metas'][0].data[0] - imgs = tensor2imgs(img_tensor, **img_metas[0]['img_norm_cfg']) + img_tensor = data["img"][0] + img_metas = data["img_metas"][0].data[0] + imgs = tensor2imgs(img_tensor, **img_metas[0]["img_norm_cfg"]) assert len(imgs) == len(img_metas) for img, img_meta in zip(imgs, img_metas): - h, w, _ = img_meta['img_shape'] + h, w, _ = img_meta["img_shape"] img_show = img[:h, :w, :] - ori_h, ori_w = img_meta['ori_shape'][:-1] + ori_h, ori_w = img_meta["ori_shape"][:-1] img_show = mmcv.imresize(img_show, (ori_w, ori_h)) if out_dir: - out_file = osp.join(out_dir, img_meta['ori_filename']) + out_file = osp.join(out_dir, img_meta["ori_filename"]) else: out_file = None model.module.show_result( - img_show, - result, - palette=dataset.PALETTE, - show=show, - out_file=out_file, - opacity=opacity) + img_show, result, palette=dataset.PALETTE, show=show, out_file=out_file, opacity=opacity + ) if isinstance(result, list): if efficient_test: @@ -103,11 +93,7 @@ def single_gpu_test(model, return results -def multi_gpu_test(model, - data_loader, - tmpdir=None, - gpu_collect=False, - efficient_test=False): +def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False, efficient_test=False): """Test model with multiple gpus. This method tests model with multiple gpus and collects the results @@ -149,7 +135,7 @@ def multi_gpu_test(model, results.append(result) if rank == 0: - batch_size = data['img'][0].size(0) + batch_size = data["img"][0].size(0) for _ in range(batch_size * world_size): prog_bar.update() @@ -168,21 +154,17 @@ def collect_results_cpu(result_part, size, tmpdir=None): if tmpdir is None: MAX_LEN = 512 # 32 is whitespace - dir_tensor = torch.full((MAX_LEN, ), - 32, - dtype=torch.uint8, - device='cuda') + dir_tensor = torch.full((MAX_LEN,), 32, dtype=torch.uint8, device="cuda") if rank == 0: tmpdir = tempfile.mkdtemp() - tmpdir = torch.tensor( - bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda') - dir_tensor[:len(tmpdir)] = tmpdir + tmpdir = torch.tensor(bytearray(tmpdir.encode()), dtype=torch.uint8, device="cuda") + dir_tensor[: len(tmpdir)] = tmpdir dist.broadcast(dir_tensor, 0) tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip() else: mmcv.mkdir_or_exist(tmpdir) # dump the part result to the dir - mmcv.dump(result_part, osp.join(tmpdir, 'part_{}.pkl'.format(rank))) + mmcv.dump(result_part, osp.join(tmpdir, "part_{}.pkl".format(rank))) dist.barrier() # collect all parts if rank != 0: @@ -191,7 +173,7 @@ def collect_results_cpu(result_part, size, tmpdir=None): # load results of all parts from tmp dir part_list = [] for i in range(world_size): - part_file = osp.join(tmpdir, 'part_{}.pkl'.format(i)) + part_file = osp.join(tmpdir, "part_{}.pkl".format(i)) part_list.append(mmcv.load(part_file)) # sort the results ordered_results = [] @@ -208,27 +190,23 @@ def collect_results_gpu(result_part, size): """Collect results with GPU.""" rank, world_size = get_dist_info() # dump result part to tensor with pickle - part_tensor = torch.tensor( - bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda') + part_tensor = torch.tensor(bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device="cuda") # gather all result part tensor shape - shape_tensor = torch.tensor(part_tensor.shape, device='cuda') + shape_tensor = torch.tensor(part_tensor.shape, device="cuda") shape_list = [shape_tensor.clone() for _ in range(world_size)] dist.all_gather(shape_list, shape_tensor) # padding result part tensor to max length shape_max = torch.tensor(shape_list).max() - part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda') - part_send[:shape_tensor[0]] = part_tensor - part_recv_list = [ - part_tensor.new_zeros(shape_max) for _ in range(world_size) - ] + part_send = torch.zeros(shape_max, dtype=torch.uint8, device="cuda") + part_send[: shape_tensor[0]] = part_tensor + part_recv_list = [part_tensor.new_zeros(shape_max) for _ in range(world_size)] # gather all result part dist.all_gather(part_recv_list, part_send) if rank == 0: part_list = [] for recv, shape in zip(part_recv_list, shape_list): - part_list.append( - pickle.loads(recv[:shape[0]].cpu().numpy().tobytes())) + part_list.append(pickle.loads(recv[: shape[0]].cpu().numpy().tobytes())) # sort the results ordered_results = [] for res in zip(*part_list): diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/train.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/train.py index ec74f20a..f10601a1 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/train.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/apis/train.py @@ -3,12 +3,15 @@ import numpy as np import torch -from comfy_controlnet_preprocessors.uniformer.mmcv.parallel import MMDataParallel, MMDistributedDataParallel -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import build_optimizer, build_runner +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.parallel import ( + MMDataParallel, + MMDistributedDataParallel, +) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import build_optimizer, build_runner -from comfy_controlnet_preprocessors.uniformer.mmseg.core import DistEvalHook, EvalHook -from comfy_controlnet_preprocessors.uniformer.mmseg.datasets import build_dataloader, build_dataset -from comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.core import DistEvalHook, EvalHook +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.datasets import build_dataloader, build_dataset +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger def set_random_seed(seed, deterministic=False): @@ -30,13 +33,7 @@ def set_random_seed(seed, deterministic=False): torch.backends.cudnn.benchmark = False -def train_segmentor(model, - dataset, - cfg, - distributed=False, - validate=False, - timestamp=None, - meta=None): +def train_segmentor(model, dataset, cfg, distributed=False, validate=False, timestamp=None, meta=None): """Launch segmentor training.""" logger = get_root_logger(cfg.log_level) @@ -51,46 +48,45 @@ def train_segmentor(model, len(cfg.gpu_ids), dist=distributed, seed=cfg.seed, - drop_last=True) for ds in dataset + drop_last=True, + ) + for ds in dataset ] # put model on gpus if distributed: - find_unused_parameters = cfg.get('find_unused_parameters', False) + find_unused_parameters = cfg.get("find_unused_parameters", False) # Sets the `find_unused_parameters` parameter in # torch.nn.parallel.DistributedDataParallel model = MMDistributedDataParallel( model.cuda(), device_ids=[torch.cuda.current_device()], broadcast_buffers=False, - find_unused_parameters=find_unused_parameters) + find_unused_parameters=find_unused_parameters, + ) else: - model = MMDataParallel( - model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids) + model = MMDataParallel(model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids) # build runner optimizer = build_optimizer(model, cfg.optimizer) - if cfg.get('runner') is None: - cfg.runner = {'type': 'IterBasedRunner', 'max_iters': cfg.total_iters} + if cfg.get("runner") is None: + cfg.runner = {"type": "IterBasedRunner", "max_iters": cfg.total_iters} warnings.warn( - 'config is now expected to have a `runner` section, ' - 'please set `runner` in your config.', UserWarning) + "config is now expected to have a `runner` section, " "please set `runner` in your config.", UserWarning + ) runner = build_runner( cfg.runner, default_args=dict( - model=model, - batch_processor=None, - optimizer=optimizer, - work_dir=cfg.work_dir, - logger=logger, - meta=meta)) + model=model, batch_processor=None, optimizer=optimizer, work_dir=cfg.work_dir, logger=logger, meta=meta + ), + ) # register hooks - runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config, - cfg.checkpoint_config, cfg.log_config, - cfg.get('momentum_config', None)) + runner.register_training_hooks( + cfg.lr_config, cfg.optimizer_config, cfg.checkpoint_config, cfg.log_config, cfg.get("momentum_config", None) + ) # an ugly walkaround to make the .log and .log.json filenames the same runner.timestamp = timestamp @@ -99,15 +95,12 @@ def train_segmentor(model, if validate: val_dataset = build_dataset(cfg.data.val, dict(test_mode=True)) val_dataloader = build_dataloader( - val_dataset, - samples_per_gpu=1, - workers_per_gpu=cfg.data.workers_per_gpu, - dist=distributed, - shuffle=False) - eval_cfg = cfg.get('evaluation', {}) - eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner' + val_dataset, samples_per_gpu=1, workers_per_gpu=cfg.data.workers_per_gpu, dist=distributed, shuffle=False + ) + eval_cfg = cfg.get("evaluation", {}) + eval_cfg["by_epoch"] = cfg.runner["type"] != "IterBasedRunner" eval_hook = DistEvalHook if distributed else EvalHook - runner.register_hook(eval_hook(val_dataloader, **eval_cfg), priority='LOW') + runner.register_hook(eval_hook(val_dataloader, **eval_cfg), priority="LOW") if cfg.resume_from: runner.resume(cfg.resume_from) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/class_names.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/class_names.py index 56c2a8f5..2ce0424d 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/class_names.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/class_names.py @@ -1,120 +1,426 @@ -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv def cityscapes_classes(): """Cityscapes class names for external use.""" return [ - 'road', 'sidewalk', 'building', 'wall', 'fence', 'pole', - 'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky', - 'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', - 'bicycle' + "road", + "sidewalk", + "building", + "wall", + "fence", + "pole", + "traffic light", + "traffic sign", + "vegetation", + "terrain", + "sky", + "person", + "rider", + "car", + "truck", + "bus", + "train", + "motorcycle", + "bicycle", ] def ade_classes(): """ADE20K class names for external use.""" return [ - 'wall', 'building', 'sky', 'floor', 'tree', 'ceiling', 'road', 'bed ', - 'windowpane', 'grass', 'cabinet', 'sidewalk', 'person', 'earth', - 'door', 'table', 'mountain', 'plant', 'curtain', 'chair', 'car', - 'water', 'painting', 'sofa', 'shelf', 'house', 'sea', 'mirror', 'rug', - 'field', 'armchair', 'seat', 'fence', 'desk', 'rock', 'wardrobe', - 'lamp', 'bathtub', 'railing', 'cushion', 'base', 'box', 'column', - 'signboard', 'chest of drawers', 'counter', 'sand', 'sink', - 'skyscraper', 'fireplace', 'refrigerator', 'grandstand', 'path', - 'stairs', 'runway', 'case', 'pool table', 'pillow', 'screen door', - 'stairway', 'river', 'bridge', 'bookcase', 'blind', 'coffee table', - 'toilet', 'flower', 'book', 'hill', 'bench', 'countertop', 'stove', - 'palm', 'kitchen island', 'computer', 'swivel chair', 'boat', 'bar', - 'arcade machine', 'hovel', 'bus', 'towel', 'light', 'truck', 'tower', - 'chandelier', 'awning', 'streetlight', 'booth', 'television receiver', - 'airplane', 'dirt track', 'apparel', 'pole', 'land', 'bannister', - 'escalator', 'ottoman', 'bottle', 'buffet', 'poster', 'stage', 'van', - 'ship', 'fountain', 'conveyer belt', 'canopy', 'washer', 'plaything', - 'swimming pool', 'stool', 'barrel', 'basket', 'waterfall', 'tent', - 'bag', 'minibike', 'cradle', 'oven', 'ball', 'food', 'step', 'tank', - 'trade name', 'microwave', 'pot', 'animal', 'bicycle', 'lake', - 'dishwasher', 'screen', 'blanket', 'sculpture', 'hood', 'sconce', - 'vase', 'traffic light', 'tray', 'ashcan', 'fan', 'pier', 'crt screen', - 'plate', 'monitor', 'bulletin board', 'shower', 'radiator', 'glass', - 'clock', 'flag' + "wall", + "building", + "sky", + "floor", + "tree", + "ceiling", + "road", + "bed ", + "windowpane", + "grass", + "cabinet", + "sidewalk", + "person", + "earth", + "door", + "table", + "mountain", + "plant", + "curtain", + "chair", + "car", + "water", + "painting", + "sofa", + "shelf", + "house", + "sea", + "mirror", + "rug", + "field", + "armchair", + "seat", + "fence", + "desk", + "rock", + "wardrobe", + "lamp", + "bathtub", + "railing", + "cushion", + "base", + "box", + "column", + "signboard", + "chest of drawers", + "counter", + "sand", + "sink", + "skyscraper", + "fireplace", + "refrigerator", + "grandstand", + "path", + "stairs", + "runway", + "case", + "pool table", + "pillow", + "screen door", + "stairway", + "river", + "bridge", + "bookcase", + "blind", + "coffee table", + "toilet", + "flower", + "book", + "hill", + "bench", + "countertop", + "stove", + "palm", + "kitchen island", + "computer", + "swivel chair", + "boat", + "bar", + "arcade machine", + "hovel", + "bus", + "towel", + "light", + "truck", + "tower", + "chandelier", + "awning", + "streetlight", + "booth", + "television receiver", + "airplane", + "dirt track", + "apparel", + "pole", + "land", + "bannister", + "escalator", + "ottoman", + "bottle", + "buffet", + "poster", + "stage", + "van", + "ship", + "fountain", + "conveyer belt", + "canopy", + "washer", + "plaything", + "swimming pool", + "stool", + "barrel", + "basket", + "waterfall", + "tent", + "bag", + "minibike", + "cradle", + "oven", + "ball", + "food", + "step", + "tank", + "trade name", + "microwave", + "pot", + "animal", + "bicycle", + "lake", + "dishwasher", + "screen", + "blanket", + "sculpture", + "hood", + "sconce", + "vase", + "traffic light", + "tray", + "ashcan", + "fan", + "pier", + "crt screen", + "plate", + "monitor", + "bulletin board", + "shower", + "radiator", + "glass", + "clock", + "flag", ] def voc_classes(): """Pascal VOC class names for external use.""" return [ - 'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', - 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', - 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', - 'tvmonitor' + "background", + "aeroplane", + "bicycle", + "bird", + "boat", + "bottle", + "bus", + "car", + "cat", + "chair", + "cow", + "diningtable", + "dog", + "horse", + "motorbike", + "person", + "pottedplant", + "sheep", + "sofa", + "train", + "tvmonitor", ] def cityscapes_palette(): """Cityscapes palette for external use.""" - return [[128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156], - [190, 153, 153], [153, 153, 153], [250, 170, 30], [220, 220, 0], - [107, 142, 35], [152, 251, 152], [70, 130, 180], [220, 20, 60], - [255, 0, 0], [0, 0, 142], [0, 0, 70], [0, 60, 100], [0, 80, 100], - [0, 0, 230], [119, 11, 32]] + return [ + [128, 64, 128], + [244, 35, 232], + [70, 70, 70], + [102, 102, 156], + [190, 153, 153], + [153, 153, 153], + [250, 170, 30], + [220, 220, 0], + [107, 142, 35], + [152, 251, 152], + [70, 130, 180], + [220, 20, 60], + [255, 0, 0], + [0, 0, 142], + [0, 0, 70], + [0, 60, 100], + [0, 80, 100], + [0, 0, 230], + [119, 11, 32], + ] def ade_palette(): """ADE20K palette for external use.""" - return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50], - [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255], - [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7], - [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82], - [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3], - [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255], - [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220], - [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224], - [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255], - [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7], - [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153], - [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255], - [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0], - [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255], - [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255], - [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255], - [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0], - [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0], - [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255], - [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255], - [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20], - [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255], - [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255], - [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255], - [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0], - [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0], - [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255], - [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112], - [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160], - [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163], - [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0], - [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0], - [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255], - [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204], - [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255], - [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255], - [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194], - [102, 255, 0], [92, 0, 255]] + return [ + [120, 120, 120], + [180, 120, 120], + [6, 230, 230], + [80, 50, 50], + [4, 200, 3], + [120, 120, 80], + [140, 140, 140], + [204, 5, 255], + [230, 230, 230], + [4, 250, 7], + [224, 5, 255], + [235, 255, 7], + [150, 5, 61], + [120, 120, 70], + [8, 255, 51], + [255, 6, 82], + [143, 255, 140], + [204, 255, 4], + [255, 51, 7], + [204, 70, 3], + [0, 102, 200], + [61, 230, 250], + [255, 6, 51], + [11, 102, 255], + [255, 7, 71], + [255, 9, 224], + [9, 7, 230], + [220, 220, 220], + [255, 9, 92], + [112, 9, 255], + [8, 255, 214], + [7, 255, 224], + [255, 184, 6], + [10, 255, 71], + [255, 41, 10], + [7, 255, 255], + [224, 255, 8], + [102, 8, 255], + [255, 61, 6], + [255, 194, 7], + [255, 122, 8], + [0, 255, 20], + [255, 8, 41], + [255, 5, 153], + [6, 51, 255], + [235, 12, 255], + [160, 150, 20], + [0, 163, 255], + [140, 140, 140], + [250, 10, 15], + [20, 255, 0], + [31, 255, 0], + [255, 31, 0], + [255, 224, 0], + [153, 255, 0], + [0, 0, 255], + [255, 71, 0], + [0, 235, 255], + [0, 173, 255], + [31, 0, 255], + [11, 200, 200], + [255, 82, 0], + [0, 255, 245], + [0, 61, 255], + [0, 255, 112], + [0, 255, 133], + [255, 0, 0], + [255, 163, 0], + [255, 102, 0], + [194, 255, 0], + [0, 143, 255], + [51, 255, 0], + [0, 82, 255], + [0, 255, 41], + [0, 255, 173], + [10, 0, 255], + [173, 255, 0], + [0, 255, 153], + [255, 92, 0], + [255, 0, 255], + [255, 0, 245], + [255, 0, 102], + [255, 173, 0], + [255, 0, 20], + [255, 184, 184], + [0, 31, 255], + [0, 255, 61], + [0, 71, 255], + [255, 0, 204], + [0, 255, 194], + [0, 255, 82], + [0, 10, 255], + [0, 112, 255], + [51, 0, 255], + [0, 194, 255], + [0, 122, 255], + [0, 255, 163], + [255, 153, 0], + [0, 255, 10], + [255, 112, 0], + [143, 255, 0], + [82, 0, 255], + [163, 255, 0], + [255, 235, 0], + [8, 184, 170], + [133, 0, 255], + [0, 255, 92], + [184, 0, 255], + [255, 0, 31], + [0, 184, 255], + [0, 214, 255], + [255, 0, 112], + [92, 255, 0], + [0, 224, 255], + [112, 224, 255], + [70, 184, 160], + [163, 0, 255], + [153, 0, 255], + [71, 255, 0], + [255, 0, 163], + [255, 204, 0], + [255, 0, 143], + [0, 255, 235], + [133, 255, 0], + [255, 0, 235], + [245, 0, 255], + [255, 0, 122], + [255, 245, 0], + [10, 190, 212], + [214, 255, 0], + [0, 204, 255], + [20, 0, 255], + [255, 255, 0], + [0, 153, 255], + [0, 41, 255], + [0, 255, 204], + [41, 0, 255], + [41, 255, 0], + [173, 0, 255], + [0, 245, 255], + [71, 0, 255], + [122, 0, 255], + [0, 255, 184], + [0, 92, 255], + [184, 255, 0], + [0, 133, 255], + [255, 214, 0], + [25, 194, 194], + [102, 255, 0], + [92, 0, 255], + ] def voc_palette(): """Pascal VOC palette for external use.""" - return [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128], - [128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], - [192, 0, 0], [64, 128, 0], [192, 128, 0], [64, 0, 128], - [192, 0, 128], [64, 128, 128], [192, 128, 128], [0, 64, 0], - [128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64, 128]] + return [ + [0, 0, 0], + [128, 0, 0], + [0, 128, 0], + [128, 128, 0], + [0, 0, 128], + [128, 0, 128], + [0, 128, 128], + [128, 128, 128], + [64, 0, 0], + [192, 0, 0], + [64, 128, 0], + [192, 128, 0], + [64, 0, 128], + [192, 0, 128], + [64, 128, 128], + [192, 128, 128], + [0, 64, 0], + [128, 64, 0], + [0, 192, 0], + [128, 192, 0], + [0, 64, 128], + ] dataset_aliases = { - 'cityscapes': ['cityscapes'], - 'ade': ['ade', 'ade20k'], - 'voc': ['voc', 'pascal_voc', 'voc12', 'voc12aug'] + "cityscapes": ["cityscapes"], + "ade": ["ade", "ade20k"], + "voc": ["voc", "pascal_voc", "voc12", "voc12aug"], } @@ -127,11 +433,11 @@ def get_classes(dataset): if mmcv.is_str(dataset): if dataset in alias2name: - labels = eval(alias2name[dataset] + '_classes()') + labels = eval(alias2name[dataset] + "_classes()") else: - raise ValueError(f'Unrecognized dataset: {dataset}') + raise ValueError(f"Unrecognized dataset: {dataset}") else: - raise TypeError(f'dataset must a str, but got {type(dataset)}') + raise TypeError(f"dataset must a str, but got {type(dataset)}") return labels @@ -144,9 +450,9 @@ def get_palette(dataset): if mmcv.is_str(dataset): if dataset in alias2name: - labels = eval(alias2name[dataset] + '_palette()') + labels = eval(alias2name[dataset] + "_palette()") else: - raise ValueError(f'Unrecognized dataset: {dataset}') + raise ValueError(f"Unrecognized dataset: {dataset}") else: - raise TypeError(f'dataset must a str, but got {type(dataset)}') + raise TypeError(f"dataset must a str, but got {type(dataset)}") return labels diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/eval_hooks.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/eval_hooks.py index 9222c315..1d54dc1f 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/eval_hooks.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/eval_hooks.py @@ -1,7 +1,7 @@ import os.path as osp -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import DistEvalHook as _DistEvalHook -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import EvalHook as _EvalHook +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import DistEvalHook as _DistEvalHook +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import EvalHook as _EvalHook class EvalHook(_EvalHook): @@ -17,7 +17,7 @@ class EvalHook(_EvalHook): list: The prediction results. """ - greater_keys = ['mIoU', 'mAcc', 'aAcc'] + greater_keys = ["mIoU", "mAcc", "aAcc"] def __init__(self, *args, by_epoch=False, efficient_test=False, **kwargs): super().__init__(*args, by_epoch=by_epoch, **kwargs) @@ -30,13 +30,10 @@ def after_train_iter(self, runner): """ if self.by_epoch or not self.every_n_iters(runner, self.interval): return - from comfy_controlnet_preprocessors.uniformer.mmseg.apis import single_gpu_test + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.apis import single_gpu_test + runner.log_buffer.clear() - results = single_gpu_test( - runner.model, - self.dataloader, - show=False, - efficient_test=self.efficient_test) + results = single_gpu_test(runner.model, self.dataloader, show=False, efficient_test=self.efficient_test) self.evaluate(runner, results) def after_train_epoch(self, runner): @@ -46,7 +43,8 @@ def after_train_epoch(self, runner): """ if not self.by_epoch or not self.every_n_epochs(runner, self.interval): return - from comfy_controlnet_preprocessors.uniformer.mmseg.apis import single_gpu_test + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.apis import single_gpu_test + runner.log_buffer.clear() results = single_gpu_test(runner.model, self.dataloader, show=False) self.evaluate(runner, results) @@ -65,7 +63,7 @@ class DistEvalHook(_DistEvalHook): list: The prediction results. """ - greater_keys = ['mIoU', 'mAcc', 'aAcc'] + greater_keys = ["mIoU", "mAcc", "aAcc"] def __init__(self, *args, by_epoch=False, efficient_test=False, **kwargs): super().__init__(*args, by_epoch=by_epoch, **kwargs) @@ -78,16 +76,18 @@ def after_train_iter(self, runner): """ if self.by_epoch or not self.every_n_iters(runner, self.interval): return - from comfy_controlnet_preprocessors.uniformer.mmseg.apis import multi_gpu_test + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.apis import multi_gpu_test + runner.log_buffer.clear() results = multi_gpu_test( runner.model, self.dataloader, - tmpdir=osp.join(runner.work_dir, '.eval_hook'), + tmpdir=osp.join(runner.work_dir, ".eval_hook"), gpu_collect=self.gpu_collect, - efficient_test=self.efficient_test) + efficient_test=self.efficient_test, + ) if runner.rank == 0: - print('\n') + print("\n") self.evaluate(runner, results) def after_train_epoch(self, runner): @@ -97,13 +97,12 @@ def after_train_epoch(self, runner): """ if not self.by_epoch or not self.every_n_epochs(runner, self.interval): return - from comfy_controlnet_preprocessors.uniformer.mmseg.apis import multi_gpu_test + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.apis import multi_gpu_test + runner.log_buffer.clear() results = multi_gpu_test( - runner.model, - self.dataloader, - tmpdir=osp.join(runner.work_dir, '.eval_hook'), - gpu_collect=self.gpu_collect) + runner.model, self.dataloader, tmpdir=osp.join(runner.work_dir, ".eval_hook"), gpu_collect=self.gpu_collect + ) if runner.rank == 0: - print('\n') + print("\n") self.evaluate(runner, results) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/metrics.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/metrics.py index d27d9947..f49d1205 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/metrics.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/evaluation/metrics.py @@ -1,6 +1,6 @@ from collections import OrderedDict -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import numpy as np import torch @@ -17,17 +17,11 @@ def f_score(precision, recall, beta=1): Returns: [torch.tensor]: The f-score value. """ - score = (1 + beta**2) * (precision * recall) / ( - (beta**2 * precision) + recall) + score = (1 + beta**2) * (precision * recall) / ((beta**2 * precision) + recall) return score -def intersect_and_union(pred_label, - label, - num_classes, - ignore_index, - label_map=dict(), - reduce_zero_label=False): +def intersect_and_union(pred_label, label, num_classes, ignore_index, label_map=dict(), reduce_zero_label=False): """Calculate intersection and Union. Args: @@ -57,8 +51,7 @@ def intersect_and_union(pred_label, pred_label = torch.from_numpy((pred_label)) if isinstance(label, str): - label = torch.from_numpy( - mmcv.imread(label, flag='unchanged', backend='pillow')) + label = torch.from_numpy(mmcv.imread(label, flag="unchanged", backend="pillow")) else: label = torch.from_numpy(label) @@ -70,27 +63,21 @@ def intersect_and_union(pred_label, label = label - 1 label[label == 254] = 255 - mask = (label != ignore_index) + mask = label != ignore_index pred_label = pred_label[mask] label = label[mask] intersect = pred_label[pred_label == label] - area_intersect = torch.histc( - intersect.float(), bins=(num_classes), min=0, max=num_classes - 1) - area_pred_label = torch.histc( - pred_label.float(), bins=(num_classes), min=0, max=num_classes - 1) - area_label = torch.histc( - label.float(), bins=(num_classes), min=0, max=num_classes - 1) + area_intersect = torch.histc(intersect.float(), bins=(num_classes), min=0, max=num_classes - 1) + area_pred_label = torch.histc(pred_label.float(), bins=(num_classes), min=0, max=num_classes - 1) + area_label = torch.histc(label.float(), bins=(num_classes), min=0, max=num_classes - 1) area_union = area_pred_label + area_label - area_intersect return area_intersect, area_union, area_pred_label, area_label -def total_intersect_and_union(results, - gt_seg_maps, - num_classes, - ignore_index, - label_map=dict(), - reduce_zero_label=False): +def total_intersect_and_union( + results, gt_seg_maps, num_classes, ignore_index, label_map=dict(), reduce_zero_label=False +): """Calculate Total Intersection and Union. Args: @@ -113,30 +100,24 @@ def total_intersect_and_union(results, """ num_imgs = len(results) assert len(gt_seg_maps) == num_imgs - total_area_intersect = torch.zeros((num_classes, ), dtype=torch.float64) - total_area_union = torch.zeros((num_classes, ), dtype=torch.float64) - total_area_pred_label = torch.zeros((num_classes, ), dtype=torch.float64) - total_area_label = torch.zeros((num_classes, ), dtype=torch.float64) + total_area_intersect = torch.zeros((num_classes,), dtype=torch.float64) + total_area_union = torch.zeros((num_classes,), dtype=torch.float64) + total_area_pred_label = torch.zeros((num_classes,), dtype=torch.float64) + total_area_label = torch.zeros((num_classes,), dtype=torch.float64) for i in range(num_imgs): - area_intersect, area_union, area_pred_label, area_label = \ - intersect_and_union( - results[i], gt_seg_maps[i], num_classes, ignore_index, - label_map, reduce_zero_label) + area_intersect, area_union, area_pred_label, area_label = intersect_and_union( + results[i], gt_seg_maps[i], num_classes, ignore_index, label_map, reduce_zero_label + ) total_area_intersect += area_intersect total_area_union += area_union total_area_pred_label += area_pred_label total_area_label += area_label - return total_area_intersect, total_area_union, total_area_pred_label, \ - total_area_label + return total_area_intersect, total_area_union, total_area_pred_label, total_area_label -def mean_iou(results, - gt_seg_maps, - num_classes, - ignore_index, - nan_to_num=None, - label_map=dict(), - reduce_zero_label=False): +def mean_iou( + results, gt_seg_maps, num_classes, ignore_index, nan_to_num=None, label_map=dict(), reduce_zero_label=False +): """Calculate Mean Intersection and Union (mIoU) Args: @@ -162,20 +143,17 @@ def mean_iou(results, gt_seg_maps=gt_seg_maps, num_classes=num_classes, ignore_index=ignore_index, - metrics=['mIoU'], + metrics=["mIoU"], nan_to_num=nan_to_num, label_map=label_map, - reduce_zero_label=reduce_zero_label) + reduce_zero_label=reduce_zero_label, + ) return iou_result -def mean_dice(results, - gt_seg_maps, - num_classes, - ignore_index, - nan_to_num=None, - label_map=dict(), - reduce_zero_label=False): +def mean_dice( + results, gt_seg_maps, num_classes, ignore_index, nan_to_num=None, label_map=dict(), reduce_zero_label=False +): """Calculate Mean Dice (mDice) Args: @@ -202,21 +180,17 @@ def mean_dice(results, gt_seg_maps=gt_seg_maps, num_classes=num_classes, ignore_index=ignore_index, - metrics=['mDice'], + metrics=["mDice"], nan_to_num=nan_to_num, label_map=label_map, - reduce_zero_label=reduce_zero_label) + reduce_zero_label=reduce_zero_label, + ) return dice_result -def mean_fscore(results, - gt_seg_maps, - num_classes, - ignore_index, - nan_to_num=None, - label_map=dict(), - reduce_zero_label=False, - beta=1): +def mean_fscore( + results, gt_seg_maps, num_classes, ignore_index, nan_to_num=None, label_map=dict(), reduce_zero_label=False, beta=1 +): """Calculate Mean Intersection and Union (mIoU) Args: @@ -246,23 +220,26 @@ def mean_fscore(results, gt_seg_maps=gt_seg_maps, num_classes=num_classes, ignore_index=ignore_index, - metrics=['mFscore'], + metrics=["mFscore"], nan_to_num=nan_to_num, label_map=label_map, reduce_zero_label=reduce_zero_label, - beta=beta) + beta=beta, + ) return fscore_result -def eval_metrics(results, - gt_seg_maps, - num_classes, - ignore_index, - metrics=['mIoU'], - nan_to_num=None, - label_map=dict(), - reduce_zero_label=False, - beta=1): +def eval_metrics( + results, + gt_seg_maps, + num_classes, + ignore_index, + metrics=["mIoU"], + nan_to_num=None, + label_map=dict(), + reduce_zero_label=False, + beta=1, +): """Calculate evaluation metrics Args: results (list[ndarray] | list[str]): List of prediction segmentation @@ -283,44 +260,37 @@ def eval_metrics(results, """ if isinstance(metrics, str): metrics = [metrics] - allowed_metrics = ['mIoU', 'mDice', 'mFscore'] + allowed_metrics = ["mIoU", "mDice", "mFscore"] if not set(metrics).issubset(set(allowed_metrics)): - raise KeyError('metrics {} is not supported'.format(metrics)) + raise KeyError("metrics {} is not supported".format(metrics)) - total_area_intersect, total_area_union, total_area_pred_label, \ - total_area_label = total_intersect_and_union( - results, gt_seg_maps, num_classes, ignore_index, label_map, - reduce_zero_label) + total_area_intersect, total_area_union, total_area_pred_label, total_area_label = total_intersect_and_union( + results, gt_seg_maps, num_classes, ignore_index, label_map, reduce_zero_label + ) all_acc = total_area_intersect.sum() / total_area_label.sum() - ret_metrics = OrderedDict({'aAcc': all_acc}) + ret_metrics = OrderedDict({"aAcc": all_acc}) for metric in metrics: - if metric == 'mIoU': + if metric == "mIoU": iou = total_area_intersect / total_area_union acc = total_area_intersect / total_area_label - ret_metrics['IoU'] = iou - ret_metrics['Acc'] = acc - elif metric == 'mDice': - dice = 2 * total_area_intersect / ( - total_area_pred_label + total_area_label) + ret_metrics["IoU"] = iou + ret_metrics["Acc"] = acc + elif metric == "mDice": + dice = 2 * total_area_intersect / (total_area_pred_label + total_area_label) acc = total_area_intersect / total_area_label - ret_metrics['Dice'] = dice - ret_metrics['Acc'] = acc - elif metric == 'mFscore': + ret_metrics["Dice"] = dice + ret_metrics["Acc"] = acc + elif metric == "mFscore": precision = total_area_intersect / total_area_pred_label recall = total_area_intersect / total_area_label - f_value = torch.tensor( - [f_score(x[0], x[1], beta) for x in zip(precision, recall)]) - ret_metrics['Fscore'] = f_value - ret_metrics['Precision'] = precision - ret_metrics['Recall'] = recall - - ret_metrics = { - metric: value.numpy() - for metric, value in ret_metrics.items() - } + f_value = torch.tensor([f_score(x[0], x[1], beta) for x in zip(precision, recall)]) + ret_metrics["Fscore"] = f_value + ret_metrics["Precision"] = precision + ret_metrics["Recall"] = recall + + ret_metrics = {metric: value.numpy() for metric, value in ret_metrics.items()} if nan_to_num is not None: - ret_metrics = OrderedDict({ - metric: np.nan_to_num(metric_value, nan=nan_to_num) - for metric, metric_value in ret_metrics.items() - }) + ret_metrics = OrderedDict( + {metric: np.nan_to_num(metric_value, nan=nan_to_num) for metric, metric_value in ret_metrics.items()} + ) return ret_metrics diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/seg/builder.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/seg/builder.py index ab9fc79e..75690740 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/seg/builder.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/core/seg/builder.py @@ -1,5 +1,5 @@ import sys -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry, build_from_cfg +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry, build_from_cfg PIXEL_SAMPLERS = Registry("pixel sampler", sys.modules[__name__]) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/builder.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/builder.py index 25e6d0ad..14f449aa 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/builder.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/builder.py @@ -5,10 +5,13 @@ import sys import numpy as np -from comfy_controlnet_preprocessors.uniformer.mmcv.parallel import collate -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import get_dist_info -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry, build_from_cfg -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import DataLoader, PoolDataLoader +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.parallel import collate +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import get_dist_info +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry, build_from_cfg +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import ( + DataLoader, + PoolDataLoader, +) from torch.utils.data import DistributedSampler import torch diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/cityscapes.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/cityscapes.py index 4f3ca21a..e3ddc994 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/cityscapes.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/cityscapes.py @@ -1,9 +1,9 @@ import os.path as osp import tempfile -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import numpy as np -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import print_log +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import print_log from PIL import Image from .builder import DATASETS @@ -18,22 +18,54 @@ class CityscapesDataset(CustomDataset): fixed to '_gtFine_labelTrainIds.png' for Cityscapes dataset. """ - CLASSES = ('road', 'sidewalk', 'building', 'wall', 'fence', 'pole', - 'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky', - 'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', - 'bicycle') - - PALETTE = [[128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156], - [190, 153, 153], [153, 153, 153], [250, 170, 30], [220, 220, 0], - [107, 142, 35], [152, 251, 152], [70, 130, 180], [220, 20, 60], - [255, 0, 0], [0, 0, 142], [0, 0, 70], [0, 60, 100], - [0, 80, 100], [0, 0, 230], [119, 11, 32]] + CLASSES = ( + "road", + "sidewalk", + "building", + "wall", + "fence", + "pole", + "traffic light", + "traffic sign", + "vegetation", + "terrain", + "sky", + "person", + "rider", + "car", + "truck", + "bus", + "train", + "motorcycle", + "bicycle", + ) + + PALETTE = [ + [128, 64, 128], + [244, 35, 232], + [70, 70, 70], + [102, 102, 156], + [190, 153, 153], + [153, 153, 153], + [250, 170, 30], + [220, 220, 0], + [107, 142, 35], + [152, 251, 152], + [70, 130, 180], + [220, 20, 60], + [255, 0, 0], + [0, 0, 142], + [0, 0, 70], + [0, 60, 100], + [0, 80, 100], + [0, 0, 230], + [119, 11, 32], + ] def __init__(self, **kwargs): super(CityscapesDataset, self).__init__( - img_suffix='_leftImg8bit.png', - seg_map_suffix='_gtFine_labelTrainIds.png', - **kwargs) + img_suffix="_leftImg8bit.png", seg_map_suffix="_gtFine_labelTrainIds.png", **kwargs + ) @staticmethod def _convert_to_label_id(result): @@ -41,6 +73,7 @@ def _convert_to_label_id(result): if isinstance(result, str): result = np.load(result) import cityscapesscripts.helpers.labels as CSLabels + result_copy = result.copy() for trainId, label in CSLabels.trainId2label.items(): result_copy[result == trainId] = label.id @@ -70,13 +103,14 @@ def results2img(self, results, imgfile_prefix, to_label_id): result = results[idx] if to_label_id: result = self._convert_to_label_id(result) - filename = self.img_infos[idx]['filename'] + filename = self.img_infos[idx]["filename"] basename = osp.splitext(osp.basename(filename))[0] - png_filename = osp.join(imgfile_prefix, f'{basename}.png') + png_filename = osp.join(imgfile_prefix, f"{basename}.png") - output = Image.fromarray(result.astype(np.uint8)).convert('P') + output = Image.fromarray(result.astype(np.uint8)).convert("P") import cityscapesscripts.helpers.labels as CSLabels + palette = np.zeros((len(CSLabels.id2label), 3), dtype=np.uint8) for label_id, label in CSLabels.id2label.items(): palette[label_id] = label.color @@ -107,10 +141,10 @@ def format_results(self, results, imgfile_prefix=None, to_label_id=True): for saving json/png files when img_prefix is not specified. """ - assert isinstance(results, list), 'results must be a list' + assert isinstance(results, list), "results must be a list" assert len(results) == len(self), ( - 'The length of results is not equal to the dataset len: ' - f'{len(results)} != {len(self)}') + "The length of results is not equal to the dataset len: " f"{len(results)} != {len(self)}" + ) if imgfile_prefix is None: tmp_dir = tempfile.TemporaryDirectory() @@ -121,12 +155,7 @@ def format_results(self, results, imgfile_prefix=None, to_label_id=True): return result_files, tmp_dir - def evaluate(self, - results, - metric='mIoU', - logger=None, - imgfile_prefix=None, - efficient_test=False): + def evaluate(self, results, metric="mIoU", logger=None, imgfile_prefix=None, efficient_test=False): """Evaluation in Cityscapes/default protocol. Args: @@ -150,14 +179,11 @@ def evaluate(self, eval_results = dict() metrics = metric.copy() if isinstance(metric, list) else [metric] - if 'cityscapes' in metrics: - eval_results.update( - self._evaluate_cityscapes(results, logger, imgfile_prefix)) - metrics.remove('cityscapes') + if "cityscapes" in metrics: + eval_results.update(self._evaluate_cityscapes(results, logger, imgfile_prefix)) + metrics.remove("cityscapes") if len(metrics) > 0: - eval_results.update( - super(CityscapesDataset, - self).evaluate(results, metrics, logger, efficient_test)) + eval_results.update(super(CityscapesDataset, self).evaluate(results, metrics, logger, efficient_test)) return eval_results @@ -176,11 +202,10 @@ def _evaluate_cityscapes(self, results, logger, imgfile_prefix): try: import cityscapesscripts.evaluation.evalPixelLevelSemanticLabeling as CSEval # noqa except ImportError: - raise ImportError('Please run "pip install cityscapesscripts" to ' - 'install cityscapesscripts first.') - msg = 'Evaluating in Cityscapes style' + raise ImportError('Please run "pip install cityscapesscripts" to ' "install cityscapesscripts first.") + msg = "Evaluating in Cityscapes style" if logger is None: - msg = '\n' + msg + msg = "\n" + msg print_log(msg, logger=logger) result_files, tmp_dir = self.format_results(results, imgfile_prefix) @@ -191,7 +216,7 @@ def _evaluate_cityscapes(self, results, logger, imgfile_prefix): result_dir = tmp_dir.name eval_results = dict() - print_log(f'Evaluating results under {result_dir} ...', logger=logger) + print_log(f"Evaluating results under {result_dir} ...", logger=logger) CSEval.args.evalInstLevelScore = True CSEval.args.predictionPath = osp.abspath(result_dir) @@ -203,13 +228,11 @@ def _evaluate_cityscapes(self, results, logger, imgfile_prefix): # when evaluating with official cityscapesscripts, # **_gtFine_labelIds.png is used - for seg_map in mmcv.scandir( - self.ann_dir, 'gtFine_labelIds.png', recursive=True): + for seg_map in mmcv.scandir(self.ann_dir, "gtFine_labelIds.png", recursive=True): seg_map_list.append(osp.join(self.ann_dir, seg_map)) pred_list.append(CSEval.getPrediction(CSEval.args, seg_map)) - eval_results.update( - CSEval.evaluateImgLists(pred_list, seg_map_list, CSEval.args)) + eval_results.update(CSEval.evaluateImgLists(pred_list, seg_map_list, CSEval.args)) if tmp_dir is not None: tmp_dir.cleanup() diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/custom.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/custom.py index 97d6467a..32638962 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/custom.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/custom.py @@ -3,13 +3,13 @@ from collections import OrderedDict from functools import reduce -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import numpy as np -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import print_log +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import print_log from torch.utils.data import Dataset -from comfy_controlnet_preprocessors.uniformer.mmseg.core import eval_metrics -from comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.core import eval_metrics +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger from .builder import DATASETS from .pipelines import Compose @@ -71,19 +71,21 @@ class CustomDataset(Dataset): PALETTE = None - def __init__(self, - pipeline, - img_dir, - img_suffix='.jpg', - ann_dir=None, - seg_map_suffix='.png', - split=None, - data_root=None, - test_mode=False, - ignore_index=255, - reduce_zero_label=False, - classes=None, - palette=None): + def __init__( + self, + pipeline, + img_dir, + img_suffix=".jpg", + ann_dir=None, + seg_map_suffix=".png", + split=None, + data_root=None, + test_mode=False, + ignore_index=255, + reduce_zero_label=False, + classes=None, + palette=None, + ): self.pipeline = Compose(pipeline) self.img_dir = img_dir self.img_suffix = img_suffix @@ -95,8 +97,7 @@ def __init__(self, self.ignore_index = ignore_index self.reduce_zero_label = reduce_zero_label self.label_map = None - self.CLASSES, self.PALETTE = self.get_classes_and_palette( - classes, palette) + self.CLASSES, self.PALETTE = self.get_classes_and_palette(classes, palette) # join paths if data_root is specified if self.data_root is not None: @@ -108,16 +109,15 @@ def __init__(self, self.split = osp.join(self.data_root, self.split) # load annotations - self.img_infos = self.load_annotations(self.img_dir, self.img_suffix, - self.ann_dir, - self.seg_map_suffix, self.split) + self.img_infos = self.load_annotations( + self.img_dir, self.img_suffix, self.ann_dir, self.seg_map_suffix, self.split + ) def __len__(self): """Total number of samples of data.""" return len(self.img_infos) - def load_annotations(self, img_dir, img_suffix, ann_dir, seg_map_suffix, - split): + def load_annotations(self, img_dir, img_suffix, ann_dir, seg_map_suffix, split): """Load annotation from directory. Args: @@ -141,17 +141,17 @@ def load_annotations(self, img_dir, img_suffix, ann_dir, seg_map_suffix, img_info = dict(filename=img_name + img_suffix) if ann_dir is not None: seg_map = img_name + seg_map_suffix - img_info['ann'] = dict(seg_map=seg_map) + img_info["ann"] = dict(seg_map=seg_map) img_infos.append(img_info) else: for img in mmcv.scandir(img_dir, img_suffix, recursive=True): img_info = dict(filename=img) if ann_dir is not None: seg_map = img.replace(img_suffix, seg_map_suffix) - img_info['ann'] = dict(seg_map=seg_map) + img_info["ann"] = dict(seg_map=seg_map) img_infos.append(img_info) - print_log(f'Loaded {len(img_infos)} images', logger=get_root_logger()) + print_log(f"Loaded {len(img_infos)} images", logger=get_root_logger()) return img_infos def get_ann_info(self, idx): @@ -164,15 +164,15 @@ def get_ann_info(self, idx): dict: Annotation info of specified index. """ - return self.img_infos[idx]['ann'] + return self.img_infos[idx]["ann"] def pre_pipeline(self, results): """Prepare results dict for pipeline.""" - results['seg_fields'] = [] - results['img_prefix'] = self.img_dir - results['seg_prefix'] = self.ann_dir + results["seg_fields"] = [] + results["img_prefix"] = self.img_dir + results["seg_prefix"] = self.ann_dir if self.custom_classes: - results['label_map'] = self.label_map + results["label_map"] = self.label_map def __getitem__(self, idx): """Get training/test data after pipeline. @@ -230,12 +230,11 @@ def get_gt_seg_maps(self, efficient_test=False): """Get ground truth segmentation maps for evaluation.""" gt_seg_maps = [] for img_info in self.img_infos: - seg_map = osp.join(self.ann_dir, img_info['ann']['seg_map']) + seg_map = osp.join(self.ann_dir, img_info["ann"]["seg_map"]) if efficient_test: gt_seg_map = seg_map else: - gt_seg_map = mmcv.imread( - seg_map, flag='unchanged', backend='pillow') + gt_seg_map = mmcv.imread(seg_map, flag="unchanged", backend="pillow") gt_seg_maps.append(gt_seg_map) return gt_seg_maps @@ -263,11 +262,11 @@ def get_classes_and_palette(self, classes=None, palette=None): elif isinstance(classes, (tuple, list)): class_names = classes else: - raise ValueError(f'Unsupported type {type(classes)} of classes.') + raise ValueError(f"Unsupported type {type(classes)} of classes.") if self.CLASSES: if not set(classes).issubset(self.CLASSES): - raise ValueError('classes is not a subset of CLASSES.') + raise ValueError("classes is not a subset of CLASSES.") # dictionary, its keys are the old label ids and its values # are the new label ids. @@ -284,12 +283,10 @@ def get_classes_and_palette(self, classes=None, palette=None): return class_names, palette def get_palette_for_custom_classes(self, class_names, palette=None): - if self.label_map is not None: # return subset of palette palette = [] - for old_id, new_id in sorted( - self.label_map.items(), key=lambda x: x[1]): + for old_id, new_id in sorted(self.label_map.items(), key=lambda x: x[1]): if new_id != -1: palette.append(self.PALETTE[old_id]) palette = type(self.PALETTE)(palette) @@ -302,12 +299,7 @@ def get_palette_for_custom_classes(self, class_names, palette=None): return palette - def evaluate(self, - results, - metric='mIoU', - logger=None, - efficient_test=False, - **kwargs): + def evaluate(self, results, metric="mIoU", logger=None, efficient_test=False, **kwargs): """Evaluate the dataset. Args: @@ -323,14 +315,13 @@ def evaluate(self, if isinstance(metric, str): metric = [metric] - allowed_metrics = ['mIoU', 'mDice', 'mFscore'] + allowed_metrics = ["mIoU", "mDice", "mFscore"] if not set(metric).issubset(set(allowed_metrics)): - raise KeyError('metric {} is not supported'.format(metric)) + raise KeyError("metric {} is not supported".format(metric)) eval_results = {} gt_seg_maps = self.get_gt_seg_maps(efficient_test) if self.CLASSES is None: - num_classes = len( - reduce(np.union1d, [np.unique(_) for _ in gt_seg_maps])) + num_classes = len(reduce(np.union1d, [np.unique(_) for _ in gt_seg_maps])) else: num_classes = len(self.CLASSES) ret_metrics = eval_metrics( @@ -340,7 +331,8 @@ def evaluate(self, self.ignore_index, metric, label_map=self.label_map, - reduce_zero_label=self.reduce_zero_label) + reduce_zero_label=self.reduce_zero_label, + ) if self.CLASSES is None: class_names = tuple(range(num_classes)) @@ -348,33 +340,31 @@ def evaluate(self, class_names = self.CLASSES # summary table - ret_metrics_summary = OrderedDict({ - ret_metric: np.round(np.nanmean(ret_metric_value) * 100, 2) - for ret_metric, ret_metric_value in ret_metrics.items() - }) + ret_metrics_summary = OrderedDict( + { + ret_metric: np.round(np.nanmean(ret_metric_value) * 100, 2) + for ret_metric, ret_metric_value in ret_metrics.items() + } + ) # each class table - ret_metrics.pop('aAcc', None) - ret_metrics_class = OrderedDict({ - ret_metric: np.round(ret_metric_value * 100, 2) - for ret_metric, ret_metric_value in ret_metrics.items() - }) - ret_metrics_class.update({'Class': class_names}) - ret_metrics_class.move_to_end('Class', last=False) + ret_metrics.pop("aAcc", None) + ret_metrics_class = OrderedDict( + {ret_metric: np.round(ret_metric_value * 100, 2) for ret_metric, ret_metric_value in ret_metrics.items()} + ) + ret_metrics_class.update({"Class": class_names}) + ret_metrics_class.move_to_end("Class", last=False) # each metric dict for key, value in ret_metrics_summary.items(): - if key == 'aAcc': + if key == "aAcc": eval_results[key] = value / 100.0 else: - eval_results['m' + key] = value / 100.0 + eval_results["m" + key] = value / 100.0 - ret_metrics_class.pop('Class', None) + ret_metrics_class.pop("Class", None) for key, value in ret_metrics_class.items(): - eval_results.update({ - key + '.' + str(name): value[idx] / 100.0 - for idx, name in enumerate(class_names) - }) + eval_results.update({key + "." + str(name): value[idx] / 100.0 for idx, name in enumerate(class_names)}) if mmcv.is_list_of(results, str): for file_name in results: diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/compose.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/compose.py index 4669173e..8fc45401 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/compose.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/compose.py @@ -1,6 +1,6 @@ import collections -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import build_from_cfg +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import build_from_cfg from ..builder import PIPELINES @@ -24,7 +24,7 @@ def __init__(self, transforms): elif callable(transform): self.transforms.append(transform) else: - raise TypeError('transform must be callable or a dict') + raise TypeError("transform must be callable or a dict") def __call__(self, data): """Call function to apply transforms sequentially. @@ -43,9 +43,9 @@ def __call__(self, data): return data def __repr__(self): - format_string = self.__class__.__name__ + '(' + format_string = self.__class__.__name__ + "(" for t in self.transforms: - format_string += '\n' - format_string += f' {t}' - format_string += '\n)' + format_string += "\n" + format_string += f" {t}" + format_string += "\n)" return format_string diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/formating.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/formating.py index 37f261b9..cc549894 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/formating.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/formating.py @@ -1,9 +1,9 @@ from collections.abc import Sequence -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import numpy as np import torch -from comfy_controlnet_preprocessors.uniformer.mmcv.parallel import DataContainer as DC +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.parallel import DataContainer as DC from ..builder import PIPELINES @@ -30,7 +30,7 @@ def to_tensor(data): elif isinstance(data, float): return torch.FloatTensor([data]) else: - raise TypeError(f'type {type(data)} cannot be converted to tensor.') + raise TypeError(f"type {type(data)} cannot be converted to tensor.") @PIPELINES.register_module() @@ -60,7 +60,7 @@ def __call__(self, results): return results def __repr__(self): - return self.__class__.__name__ + f'(keys={self.keys})' + return self.__class__.__name__ + f"(keys={self.keys})" @PIPELINES.register_module() @@ -98,7 +98,7 @@ def __call__(self, results): return results def __repr__(self): - return self.__class__.__name__ + f'(keys={self.keys})' + return self.__class__.__name__ + f"(keys={self.keys})" @PIPELINES.register_module() @@ -131,8 +131,7 @@ def __call__(self, results): return results def __repr__(self): - return self.__class__.__name__ + \ - f'(keys={self.keys}, order={self.order})' + return self.__class__.__name__ + f"(keys={self.keys}, order={self.order})" @PIPELINES.register_module() @@ -147,9 +146,7 @@ class ToDataContainer(object): dict(key='gt_semantic_seg'))``. """ - def __init__(self, - fields=(dict(key='img', - stack=True), dict(key='gt_semantic_seg'))): + def __init__(self, fields=(dict(key="img", stack=True), dict(key="gt_semantic_seg"))): self.fields = fields def __call__(self, results): @@ -166,12 +163,12 @@ def __call__(self, results): for field in self.fields: field = field.copy() - key = field.pop('key') + key = field.pop("key") results[key] = DC(results[key], **field) return results def __repr__(self): - return self.__class__.__name__ + f'(fields={self.fields})' + return self.__class__.__name__ + f"(fields={self.fields})" @PIPELINES.register_module() @@ -197,18 +194,17 @@ def __call__(self, results): default bundle. """ - if 'img' in results: - img = results['img'] + if "img" in results: + img = results["img"] if len(img.shape) < 3: img = np.expand_dims(img, -1) img = np.ascontiguousarray(img.transpose(2, 0, 1)) - results['img'] = DC(to_tensor(img), stack=True) - if 'gt_semantic_seg' in results: + results["img"] = DC(to_tensor(img), stack=True) + if "gt_semantic_seg" in results: # convert to long - results['gt_semantic_seg'] = DC( - to_tensor(results['gt_semantic_seg'][None, - ...].astype(np.int64)), - stack=True) + results["gt_semantic_seg"] = DC( + to_tensor(results["gt_semantic_seg"][None, ...].astype(np.int64)), stack=True + ) return results def __repr__(self): @@ -253,11 +249,21 @@ class Collect(object): 'img_norm_cfg')`` """ - def __init__(self, - keys, - meta_keys=('filename', 'ori_filename', 'ori_shape', - 'img_shape', 'pad_shape', 'scale_factor', 'flip', - 'flip_direction', 'img_norm_cfg')): + def __init__( + self, + keys, + meta_keys=( + "filename", + "ori_filename", + "ori_shape", + "img_shape", + "pad_shape", + "scale_factor", + "flip", + "flip_direction", + "img_norm_cfg", + ), + ): self.keys = keys self.meta_keys = meta_keys @@ -278,11 +284,10 @@ def __call__(self, results): img_meta = {} for key in self.meta_keys: img_meta[key] = results[key] - data['img_metas'] = DC(img_meta, cpu_only=True) + data["img_metas"] = DC(img_meta, cpu_only=True) for key in self.keys: data[key] = results[key] return data def __repr__(self): - return self.__class__.__name__ + \ - f'(keys={self.keys}, meta_keys={self.meta_keys})' + return self.__class__.__name__ + f"(keys={self.keys}, meta_keys={self.meta_keys})" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/loading.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/loading.py index 91d5f7f9..a94654a1 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/loading.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/loading.py @@ -1,6 +1,6 @@ import os.path as osp -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import numpy as np from ..builder import PIPELINES @@ -28,11 +28,9 @@ class LoadImageFromFile(object): 'cv2' """ - def __init__(self, - to_float32=False, - color_type='color', - file_client_args=dict(backend='disk'), - imdecode_backend='cv2'): + def __init__( + self, to_float32=False, color_type="color", file_client_args=dict(backend="disk"), imdecode_backend="cv2" + ): self.to_float32 = to_float32 self.color_type = color_type self.file_client_args = file_client_args.copy() @@ -52,35 +50,32 @@ def __call__(self, results): if self.file_client is None: self.file_client = mmcv.FileClient(**self.file_client_args) - if results.get('img_prefix') is not None: - filename = osp.join(results['img_prefix'], - results['img_info']['filename']) + if results.get("img_prefix") is not None: + filename = osp.join(results["img_prefix"], results["img_info"]["filename"]) else: - filename = results['img_info']['filename'] + filename = results["img_info"]["filename"] img_bytes = self.file_client.get(filename) - img = mmcv.imfrombytes( - img_bytes, flag=self.color_type, backend=self.imdecode_backend) + img = mmcv.imfrombytes(img_bytes, flag=self.color_type, backend=self.imdecode_backend) if self.to_float32: img = img.astype(np.float32) - results['filename'] = filename - results['ori_filename'] = results['img_info']['filename'] - results['img'] = img - results['img_shape'] = img.shape - results['ori_shape'] = img.shape + results["filename"] = filename + results["ori_filename"] = results["img_info"]["filename"] + results["img"] = img + results["img_shape"] = img.shape + results["ori_shape"] = img.shape # Set initial values for default meta_keys - results['pad_shape'] = img.shape - results['scale_factor'] = 1.0 + results["pad_shape"] = img.shape + results["scale_factor"] = 1.0 num_channels = 1 if len(img.shape) < 3 else img.shape[2] - results['img_norm_cfg'] = dict( - mean=np.zeros(num_channels, dtype=np.float32), - std=np.ones(num_channels, dtype=np.float32), - to_rgb=False) + results["img_norm_cfg"] = dict( + mean=np.zeros(num_channels, dtype=np.float32), std=np.ones(num_channels, dtype=np.float32), to_rgb=False + ) return results def __repr__(self): repr_str = self.__class__.__name__ - repr_str += f'(to_float32={self.to_float32},' + repr_str += f"(to_float32={self.to_float32}," repr_str += f"color_type='{self.color_type}'," repr_str += f"imdecode_backend='{self.imdecode_backend}')" return repr_str @@ -101,10 +96,7 @@ class LoadAnnotations(object): 'pillow' """ - def __init__(self, - reduce_zero_label=False, - file_client_args=dict(backend='disk'), - imdecode_backend='pillow'): + def __init__(self, reduce_zero_label=False, file_client_args=dict(backend="disk"), imdecode_backend="pillow"): self.reduce_zero_label = reduce_zero_label self.file_client_args = file_client_args.copy() self.file_client = None @@ -123,18 +115,17 @@ def __call__(self, results): if self.file_client is None: self.file_client = mmcv.FileClient(**self.file_client_args) - if results.get('seg_prefix', None) is not None: - filename = osp.join(results['seg_prefix'], - results['ann_info']['seg_map']) + if results.get("seg_prefix", None) is not None: + filename = osp.join(results["seg_prefix"], results["ann_info"]["seg_map"]) else: - filename = results['ann_info']['seg_map'] + filename = results["ann_info"]["seg_map"] img_bytes = self.file_client.get(filename) - gt_semantic_seg = mmcv.imfrombytes( - img_bytes, flag='unchanged', - backend=self.imdecode_backend).squeeze().astype(np.uint8) + gt_semantic_seg = ( + mmcv.imfrombytes(img_bytes, flag="unchanged", backend=self.imdecode_backend).squeeze().astype(np.uint8) + ) # modify if custom classes - if results.get('label_map', None) is not None: - for old_id, new_id in results['label_map'].items(): + if results.get("label_map", None) is not None: + for old_id, new_id in results["label_map"].items(): gt_semantic_seg[gt_semantic_seg == old_id] = new_id # reduce zero_label if self.reduce_zero_label: @@ -142,12 +133,12 @@ def __call__(self, results): gt_semantic_seg[gt_semantic_seg == 0] = 255 gt_semantic_seg = gt_semantic_seg - 1 gt_semantic_seg[gt_semantic_seg == 254] = 255 - results['gt_semantic_seg'] = gt_semantic_seg - results['seg_fields'].append('gt_semantic_seg') + results["gt_semantic_seg"] = gt_semantic_seg + results["seg_fields"].append("gt_semantic_seg") return results def __repr__(self): repr_str = self.__class__.__name__ - repr_str += f'(reduce_zero_label={self.reduce_zero_label},' + repr_str += f"(reduce_zero_label={self.reduce_zero_label}," repr_str += f"imdecode_backend='{self.imdecode_backend}')" return repr_str diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/test_time_aug.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/test_time_aug.py index b65e0933..099b7d3e 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/test_time_aug.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/test_time_aug.py @@ -1,6 +1,6 @@ import warnings -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv from ..builder import PIPELINES from .compose import Compose @@ -50,45 +50,31 @@ class MultiScaleFlipAug(object): It has no effect when flip == False. Default: "horizontal". """ - def __init__(self, - transforms, - img_scale, - img_ratios=None, - flip=False, - flip_direction='horizontal'): + def __init__(self, transforms, img_scale, img_ratios=None, flip=False, flip_direction="horizontal"): self.transforms = Compose(transforms) if img_ratios is not None: - img_ratios = img_ratios if isinstance(img_ratios, - list) else [img_ratios] + img_ratios = img_ratios if isinstance(img_ratios, list) else [img_ratios] assert mmcv.is_list_of(img_ratios, float) if img_scale is None: # mode 1: given img_scale=None and a range of image ratio self.img_scale = None assert mmcv.is_list_of(img_ratios, float) - elif isinstance(img_scale, tuple) and mmcv.is_list_of( - img_ratios, float): + elif isinstance(img_scale, tuple) and mmcv.is_list_of(img_ratios, float): assert len(img_scale) == 2 # mode 2: given a scale and a range of image ratio - self.img_scale = [(int(img_scale[0] * ratio), - int(img_scale[1] * ratio)) - for ratio in img_ratios] + self.img_scale = [(int(img_scale[0] * ratio), int(img_scale[1] * ratio)) for ratio in img_ratios] else: # mode 3: given multiple scales - self.img_scale = img_scale if isinstance(img_scale, - list) else [img_scale] + self.img_scale = img_scale if isinstance(img_scale, list) else [img_scale] assert mmcv.is_list_of(self.img_scale, tuple) or self.img_scale is None self.flip = flip self.img_ratios = img_ratios - self.flip_direction = flip_direction if isinstance( - flip_direction, list) else [flip_direction] + self.flip_direction = flip_direction if isinstance(flip_direction, list) else [flip_direction] assert mmcv.is_list_of(self.flip_direction, str) - if not self.flip and self.flip_direction != ['horizontal']: - warnings.warn( - 'flip_direction has no effect when flip is set to False') - if (self.flip - and not any([t['type'] == 'RandomFlip' for t in transforms])): - warnings.warn( - 'flip has no effect when RandomFlip is not in transforms') + if not self.flip and self.flip_direction != ["horizontal"]: + warnings.warn("flip_direction has no effect when flip is set to False") + if self.flip and not any([t["type"] == "RandomFlip" for t in transforms]): + warnings.warn("flip has no effect when RandomFlip is not in transforms") def __call__(self, results): """Call function to apply test time augment transforms on results. @@ -103,9 +89,8 @@ def __call__(self, results): aug_data = [] if self.img_scale is None and mmcv.is_list_of(self.img_ratios, float): - h, w = results['img'].shape[:2] - img_scale = [(int(w * ratio), int(h * ratio)) - for ratio in self.img_ratios] + h, w = results["img"].shape[:2] + img_scale = [(int(w * ratio), int(h * ratio)) for ratio in self.img_ratios] else: img_scale = self.img_scale flip_aug = [False, True] if self.flip else [False] @@ -113,9 +98,9 @@ def __call__(self, results): for flip in flip_aug: for direction in self.flip_direction: _results = results.copy() - _results['scale'] = scale - _results['flip'] = flip - _results['flip_direction'] = direction + _results["scale"] = scale + _results["flip"] = flip + _results["flip_direction"] = direction data = self.transforms(_results) aug_data.append(data) # list of dict to dict of list @@ -127,7 +112,7 @@ def __call__(self, results): def __repr__(self): repr_str = self.__class__.__name__ - repr_str += f'(transforms={self.transforms}, ' - repr_str += f'img_scale={self.img_scale}, flip={self.flip})' - repr_str += f'flip_direction={self.flip_direction}' + repr_str += f"(transforms={self.transforms}, " + repr_str += f"img_scale={self.img_scale}, flip={self.flip})" + repr_str += f"flip_direction={self.flip_direction}" return repr_str diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/transforms.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/transforms.py index 46e8eea5..bd95606d 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/transforms.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/datasets/pipelines/transforms.py @@ -1,6 +1,6 @@ -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import numpy as np -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import deprecated_api_warning, is_tuple_of +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import deprecated_api_warning, is_tuple_of from numpy import random from ..builder import PIPELINES @@ -38,11 +38,7 @@ class Resize(object): image. """ - def __init__(self, - img_scale=None, - multiscale_mode='range', - ratio_range=None, - keep_ratio=True): + def __init__(self, img_scale=None, multiscale_mode="range", ratio_range=None, keep_ratio=True): if img_scale is None: self.img_scale = None else: @@ -58,7 +54,7 @@ def __init__(self, assert self.img_scale is None or len(self.img_scale) == 1 else: # mode 3 and 4: given multiple scales or a range of scales - assert multiscale_mode in ['value', 'range'] + assert multiscale_mode in ["value", "range"] self.multiscale_mode = multiscale_mode self.ratio_range = ratio_range @@ -100,12 +96,8 @@ def random_sample(img_scales): assert mmcv.is_list_of(img_scales, tuple) and len(img_scales) == 2 img_scale_long = [max(s) for s in img_scales] img_scale_short = [min(s) for s in img_scales] - long_edge = np.random.randint( - min(img_scale_long), - max(img_scale_long) + 1) - short_edge = np.random.randint( - min(img_scale_short), - max(img_scale_short) + 1) + long_edge = np.random.randint(min(img_scale_long), max(img_scale_long) + 1) + short_edge = np.random.randint(min(img_scale_short), max(img_scale_short) + 1) img_scale = (long_edge, short_edge) return img_scale, None @@ -156,55 +148,48 @@ def _random_scale(self, results): if self.ratio_range is not None: if self.img_scale is None: - h, w = results['img'].shape[:2] - scale, scale_idx = self.random_sample_ratio((w, h), - self.ratio_range) + h, w = results["img"].shape[:2] + scale, scale_idx = self.random_sample_ratio((w, h), self.ratio_range) else: - scale, scale_idx = self.random_sample_ratio( - self.img_scale[0], self.ratio_range) + scale, scale_idx = self.random_sample_ratio(self.img_scale[0], self.ratio_range) elif len(self.img_scale) == 1: scale, scale_idx = self.img_scale[0], 0 - elif self.multiscale_mode == 'range': + elif self.multiscale_mode == "range": scale, scale_idx = self.random_sample(self.img_scale) - elif self.multiscale_mode == 'value': + elif self.multiscale_mode == "value": scale, scale_idx = self.random_select(self.img_scale) else: raise NotImplementedError - results['scale'] = scale - results['scale_idx'] = scale_idx + results["scale"] = scale + results["scale_idx"] = scale_idx def _resize_img(self, results): """Resize images with ``results['scale']``.""" if self.keep_ratio: - img, scale_factor = mmcv.imrescale( - results['img'], results['scale'], return_scale=True) + img, scale_factor = mmcv.imrescale(results["img"], results["scale"], return_scale=True) # the w_scale and h_scale has minor difference # a real fix should be done in the mmcv.imrescale in the future new_h, new_w = img.shape[:2] - h, w = results['img'].shape[:2] + h, w = results["img"].shape[:2] w_scale = new_w / w h_scale = new_h / h else: - img, w_scale, h_scale = mmcv.imresize( - results['img'], results['scale'], return_scale=True) - scale_factor = np.array([w_scale, h_scale, w_scale, h_scale], - dtype=np.float32) - results['img'] = img - results['img_shape'] = img.shape - results['pad_shape'] = img.shape # in case that there is no padding - results['scale_factor'] = scale_factor - results['keep_ratio'] = self.keep_ratio + img, w_scale, h_scale = mmcv.imresize(results["img"], results["scale"], return_scale=True) + scale_factor = np.array([w_scale, h_scale, w_scale, h_scale], dtype=np.float32) + results["img"] = img + results["img_shape"] = img.shape + results["pad_shape"] = img.shape # in case that there is no padding + results["scale_factor"] = scale_factor + results["keep_ratio"] = self.keep_ratio def _resize_seg(self, results): """Resize semantic segmentation map with ``results['scale']``.""" - for key in results.get('seg_fields', []): + for key in results.get("seg_fields", []): if self.keep_ratio: - gt_seg = mmcv.imrescale( - results[key], results['scale'], interpolation='nearest') + gt_seg = mmcv.imrescale(results[key], results["scale"], interpolation="nearest") else: - gt_seg = mmcv.imresize( - results[key], results['scale'], interpolation='nearest') + gt_seg = mmcv.imresize(results[key], results["scale"], interpolation="nearest") results[key] = gt_seg def __call__(self, results): @@ -219,7 +204,7 @@ def __call__(self, results): 'keep_ratio' keys are added into result dict. """ - if 'scale' not in results: + if "scale" not in results: self._random_scale(results) self._resize_img(results) self._resize_seg(results) @@ -227,10 +212,12 @@ def __call__(self, results): def __repr__(self): repr_str = self.__class__.__name__ - repr_str += (f'(img_scale={self.img_scale}, ' - f'multiscale_mode={self.multiscale_mode}, ' - f'ratio_range={self.ratio_range}, ' - f'keep_ratio={self.keep_ratio})') + repr_str += ( + f"(img_scale={self.img_scale}, " + f"multiscale_mode={self.multiscale_mode}, " + f"ratio_range={self.ratio_range}, " + f"keep_ratio={self.keep_ratio})" + ) return repr_str @@ -248,13 +235,13 @@ class RandomFlip(object): 'horizontal' and 'vertical'. Default: 'horizontal'. """ - @deprecated_api_warning({'flip_ratio': 'prob'}, cls_name='RandomFlip') - def __init__(self, prob=None, direction='horizontal'): + @deprecated_api_warning({"flip_ratio": "prob"}, cls_name="RandomFlip") + def __init__(self, prob=None, direction="horizontal"): self.prob = prob self.direction = direction if prob is not None: assert prob >= 0 and prob <= 1 - assert direction in ['horizontal', 'vertical'] + assert direction in ["horizontal", "vertical"] def __call__(self, results): """Call function to flip bounding boxes, masks, semantic segmentation @@ -268,25 +255,23 @@ def __call__(self, results): result dict. """ - if 'flip' not in results: + if "flip" not in results: flip = True if np.random.rand() < self.prob else False - results['flip'] = flip - if 'flip_direction' not in results: - results['flip_direction'] = self.direction - if results['flip']: + results["flip"] = flip + if "flip_direction" not in results: + results["flip_direction"] = self.direction + if results["flip"]: # flip image - results['img'] = mmcv.imflip( - results['img'], direction=results['flip_direction']) + results["img"] = mmcv.imflip(results["img"], direction=results["flip_direction"]) # flip segs - for key in results.get('seg_fields', []): + for key in results.get("seg_fields", []): # use copy() to make numpy stride positive - results[key] = mmcv.imflip( - results[key], direction=results['flip_direction']).copy() + results[key] = mmcv.imflip(results[key], direction=results["flip_direction"]).copy() return results def __repr__(self): - return self.__class__.__name__ + f'(prob={self.prob})' + return self.__class__.__name__ + f"(prob={self.prob})" @PIPELINES.register_module() @@ -305,11 +290,7 @@ class Pad(object): Default: 255. """ - def __init__(self, - size=None, - size_divisor=None, - pad_val=0, - seg_pad_val=255): + def __init__(self, size=None, size_divisor=None, pad_val=0, seg_pad_val=255): self.size = size self.size_divisor = size_divisor self.pad_val = pad_val @@ -321,23 +302,18 @@ def __init__(self, def _pad_img(self, results): """Pad images according to ``self.size``.""" if self.size is not None: - padded_img = mmcv.impad( - results['img'], shape=self.size, pad_val=self.pad_val) + padded_img = mmcv.impad(results["img"], shape=self.size, pad_val=self.pad_val) elif self.size_divisor is not None: - padded_img = mmcv.impad_to_multiple( - results['img'], self.size_divisor, pad_val=self.pad_val) - results['img'] = padded_img - results['pad_shape'] = padded_img.shape - results['pad_fixed_size'] = self.size - results['pad_size_divisor'] = self.size_divisor + padded_img = mmcv.impad_to_multiple(results["img"], self.size_divisor, pad_val=self.pad_val) + results["img"] = padded_img + results["pad_shape"] = padded_img.shape + results["pad_fixed_size"] = self.size + results["pad_size_divisor"] = self.size_divisor def _pad_seg(self, results): """Pad masks according to ``results['pad_shape']``.""" - for key in results.get('seg_fields', []): - results[key] = mmcv.impad( - results[key], - shape=results['pad_shape'][:2], - pad_val=self.seg_pad_val) + for key in results.get("seg_fields", []): + results[key] = mmcv.impad(results[key], shape=results["pad_shape"][:2], pad_val=self.seg_pad_val) def __call__(self, results): """Call function to pad images, masks, semantic segmentation maps. @@ -355,8 +331,7 @@ def __call__(self, results): def __repr__(self): repr_str = self.__class__.__name__ - repr_str += f'(size={self.size}, size_divisor={self.size_divisor}, ' \ - f'pad_val={self.pad_val})' + repr_str += f"(size={self.size}, size_divisor={self.size_divisor}, " f"pad_val={self.pad_val})" return repr_str @@ -389,16 +364,13 @@ def __call__(self, results): result dict. """ - results['img'] = mmcv.imnormalize(results['img'], self.mean, self.std, - self.to_rgb) - results['img_norm_cfg'] = dict( - mean=self.mean, std=self.std, to_rgb=self.to_rgb) + results["img"] = mmcv.imnormalize(results["img"], self.mean, self.std, self.to_rgb) + results["img_norm_cfg"] = dict(mean=self.mean, std=self.std, to_rgb=self.to_rgb) return results def __repr__(self): repr_str = self.__class__.__name__ - repr_str += f'(mean={self.mean}, std={self.std}, to_rgb=' \ - f'{self.to_rgb})' + repr_str += f"(mean={self.mean}, std={self.std}, to_rgb=" f"{self.to_rgb})" return repr_str @@ -429,7 +401,7 @@ def __call__(self, results): dict: Reranged results. """ - img = results['img'] + img = results["img"] img_min_value = np.min(img) img_max_value = np.max(img) @@ -438,13 +410,13 @@ def __call__(self, results): img = (img - img_min_value) / (img_max_value - img_min_value) # rerange to [min_value, max_value] img = img * (self.max_value - self.min_value) + self.min_value - results['img'] = img + results["img"] = img return results def __repr__(self): repr_str = self.__class__.__name__ - repr_str += f'(min_value={self.min_value}, max_value={self.max_value})' + repr_str += f"(min_value={self.min_value}, max_value={self.max_value})" return repr_str @@ -479,17 +451,16 @@ def __call__(self, results): dict: Processed results. """ - for i in range(results['img'].shape[2]): - results['img'][:, :, i] = mmcv.clahe( - np.array(results['img'][:, :, i], dtype=np.uint8), - self.clip_limit, self.tile_grid_size) + for i in range(results["img"].shape[2]): + results["img"][:, :, i] = mmcv.clahe( + np.array(results["img"][:, :, i], dtype=np.uint8), self.clip_limit, self.tile_grid_size + ) return results def __repr__(self): repr_str = self.__class__.__name__ - repr_str += f'(clip_limit={self.clip_limit}, '\ - f'tile_grid_size={self.tile_grid_size})' + repr_str += f"(clip_limit={self.clip_limit}, " f"tile_grid_size={self.tile_grid_size})" return repr_str @@ -503,7 +474,7 @@ class RandomCrop(object): occupy. """ - def __init__(self, crop_size, cat_max_ratio=1., ignore_index=255): + def __init__(self, crop_size, cat_max_ratio=1.0, ignore_index=255): assert crop_size[0] > 0 and crop_size[1] > 0 self.crop_size = crop_size self.cat_max_ratio = cat_max_ratio @@ -537,33 +508,32 @@ def __call__(self, results): updated according to crop size. """ - img = results['img'] + img = results["img"] crop_bbox = self.get_crop_bbox(img) - if self.cat_max_ratio < 1.: + if self.cat_max_ratio < 1.0: # Repeat 10 times for _ in range(10): - seg_temp = self.crop(results['gt_semantic_seg'], crop_bbox) + seg_temp = self.crop(results["gt_semantic_seg"], crop_bbox) labels, cnt = np.unique(seg_temp, return_counts=True) cnt = cnt[labels != self.ignore_index] - if len(cnt) > 1 and np.max(cnt) / np.sum( - cnt) < self.cat_max_ratio: + if len(cnt) > 1 and np.max(cnt) / np.sum(cnt) < self.cat_max_ratio: break crop_bbox = self.get_crop_bbox(img) # crop the image img = self.crop(img, crop_bbox) img_shape = img.shape - results['img'] = img - results['img_shape'] = img_shape + results["img"] = img + results["img_shape"] = img_shape # crop semantic seg - for key in results.get('seg_fields', []): + for key in results.get("seg_fields", []): results[key] = self.crop(results[key], crop_bbox) return results def __repr__(self): - return self.__class__.__name__ + f'(crop_size={self.crop_size})' + return self.__class__.__name__ + f"(crop_size={self.crop_size})" @PIPELINES.register_module() @@ -585,22 +555,15 @@ class RandomRotate(object): rotated image. Default: False """ - def __init__(self, - prob, - degree, - pad_val=0, - seg_pad_val=255, - center=None, - auto_bound=False): + def __init__(self, prob, degree, pad_val=0, seg_pad_val=255, center=None, auto_bound=False): self.prob = prob assert prob >= 0 and prob <= 1 if isinstance(degree, (float, int)): - assert degree > 0, f'degree {degree} should be positive' + assert degree > 0, f"degree {degree} should be positive" self.degree = (-degree, degree) else: self.degree = degree - assert len(self.degree) == 2, f'degree {self.degree} should be a ' \ - f'tuple of (min, max)' + assert len(self.degree) == 2, f"degree {self.degree} should be a " f"tuple of (min, max)" self.pal_val = pad_val self.seg_pad_val = seg_pad_val self.center = center @@ -620,32 +583,32 @@ def __call__(self, results): degree = np.random.uniform(min(*self.degree), max(*self.degree)) if rotate: # rotate image - results['img'] = mmcv.imrotate( - results['img'], - angle=degree, - border_value=self.pal_val, - center=self.center, - auto_bound=self.auto_bound) + results["img"] = mmcv.imrotate( + results["img"], angle=degree, border_value=self.pal_val, center=self.center, auto_bound=self.auto_bound + ) # rotate segs - for key in results.get('seg_fields', []): + for key in results.get("seg_fields", []): results[key] = mmcv.imrotate( results[key], angle=degree, border_value=self.seg_pad_val, center=self.center, auto_bound=self.auto_bound, - interpolation='nearest') + interpolation="nearest", + ) return results def __repr__(self): repr_str = self.__class__.__name__ - repr_str += f'(prob={self.prob}, ' \ - f'degree={self.degree}, ' \ - f'pad_val={self.pal_val}, ' \ - f'seg_pad_val={self.seg_pad_val}, ' \ - f'center={self.center}, ' \ - f'auto_bound={self.auto_bound})' + repr_str += ( + f"(prob={self.prob}, " + f"degree={self.degree}, " + f"pad_val={self.pal_val}, " + f"seg_pad_val={self.seg_pad_val}, " + f"center={self.center}, " + f"auto_bound={self.auto_bound})" + ) return repr_str @@ -682,7 +645,7 @@ def __call__(self, results): Returns: dict: Result dict with grayscale image. """ - img = results['img'] + img = results["img"] assert len(img.shape) == 3 assert img.shape[2] == len(self.weights) weights = np.array(self.weights).reshape((1, 1, -1)) @@ -692,15 +655,14 @@ def __call__(self, results): else: img = img.repeat(self.out_channels, axis=2) - results['img'] = img - results['img_shape'] = img.shape + results["img"] = img + results["img_shape"] = img.shape return results def __repr__(self): repr_str = self.__class__.__name__ - repr_str += f'(out_channels={self.out_channels}, ' \ - f'weights={self.weights})' + repr_str += f"(out_channels={self.out_channels}, " f"weights={self.weights})" return repr_str @@ -718,8 +680,7 @@ def __init__(self, gamma=1.0): assert gamma > 0 self.gamma = gamma inv_gamma = 1.0 / gamma - self.table = np.array([(i / 255.0)**inv_gamma * 255 - for i in np.arange(256)]).astype('uint8') + self.table = np.array([(i / 255.0) ** inv_gamma * 255 for i in np.arange(256)]).astype("uint8") def __call__(self, results): """Call function to process the image with gamma correction. @@ -731,13 +692,12 @@ def __call__(self, results): dict: Processed results. """ - results['img'] = mmcv.lut_transform( - np.array(results['img'], dtype=np.uint8), self.table) + results["img"] = mmcv.lut_transform(np.array(results["img"], dtype=np.uint8), self.table) return results def __repr__(self): - return self.__class__.__name__ + f'(gamma={self.gamma})' + return self.__class__.__name__ + f"(gamma={self.gamma})" @PIPELINES.register_module() @@ -760,14 +720,13 @@ def __call__(self, results): Returns: dict: Result dict with semantic segmentation map scaled. """ - for key in results.get('seg_fields', []): + for key in results.get("seg_fields", []): if self.scale_factor != 1: - results[key] = mmcv.imrescale( - results[key], self.scale_factor, interpolation='nearest') + results[key] = mmcv.imrescale(results[key], self.scale_factor, interpolation="nearest") return results def __repr__(self): - return self.__class__.__name__ + f'(scale_factor={self.scale_factor})' + return self.__class__.__name__ + f"(scale_factor={self.scale_factor})" @PIPELINES.register_module() @@ -791,11 +750,7 @@ class PhotoMetricDistortion(object): hue_delta (int): delta of hue. """ - def __init__(self, - brightness_delta=32, - contrast_range=(0.5, 1.5), - saturation_range=(0.5, 1.5), - hue_delta=18): + def __init__(self, brightness_delta=32, contrast_range=(0.5, 1.5), saturation_range=(0.5, 1.5), hue_delta=18): self.brightness_delta = brightness_delta self.contrast_lower, self.contrast_upper = contrast_range self.saturation_lower, self.saturation_upper = saturation_range @@ -810,18 +765,13 @@ def convert(self, img, alpha=1, beta=0): def brightness(self, img): """Brightness distortion.""" if random.randint(2): - return self.convert( - img, - beta=random.uniform(-self.brightness_delta, - self.brightness_delta)) + return self.convert(img, beta=random.uniform(-self.brightness_delta, self.brightness_delta)) return img def contrast(self, img): """Contrast distortion.""" if random.randint(2): - return self.convert( - img, - alpha=random.uniform(self.contrast_lower, self.contrast_upper)) + return self.convert(img, alpha=random.uniform(self.contrast_lower, self.contrast_upper)) return img def saturation(self, img): @@ -829,9 +779,8 @@ def saturation(self, img): if random.randint(2): img = mmcv.bgr2hsv(img) img[:, :, 1] = self.convert( - img[:, :, 1], - alpha=random.uniform(self.saturation_lower, - self.saturation_upper)) + img[:, :, 1], alpha=random.uniform(self.saturation_lower, self.saturation_upper) + ) img = mmcv.hsv2bgr(img) return img @@ -839,9 +788,7 @@ def hue(self, img): """Hue distortion.""" if random.randint(2): img = mmcv.bgr2hsv(img) - img[:, :, - 0] = (img[:, :, 0].astype(int) + - random.randint(-self.hue_delta, self.hue_delta)) % 180 + img[:, :, 0] = (img[:, :, 0].astype(int) + random.randint(-self.hue_delta, self.hue_delta)) % 180 img = mmcv.hsv2bgr(img) return img @@ -855,7 +802,7 @@ def __call__(self, results): dict: Result dict with images distorted. """ - img = results['img'] + img = results["img"] # random brightness img = self.brightness(img) @@ -875,15 +822,17 @@ def __call__(self, results): if mode == 0: img = self.contrast(img) - results['img'] = img + results["img"] = img return results def __repr__(self): repr_str = self.__class__.__name__ - repr_str += (f'(brightness_delta={self.brightness_delta}, ' - f'contrast_range=({self.contrast_lower}, ' - f'{self.contrast_upper}), ' - f'saturation_range=({self.saturation_lower}, ' - f'{self.saturation_upper}), ' - f'hue_delta={self.hue_delta})') + repr_str += ( + f"(brightness_delta={self.brightness_delta}, " + f"contrast_range=({self.contrast_lower}, " + f"{self.contrast_upper}), " + f"saturation_range=({self.saturation_lower}, " + f"{self.saturation_upper}), " + f"hue_delta={self.hue_delta})" + ) return repr_str diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/cgnet.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/cgnet.py index f19c07ae..7edd2fc9 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/cgnet.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/cgnet.py @@ -1,12 +1,17 @@ import torch import torch.nn as nn import torch.utils.checkpoint as cp -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import (ConvModule, build_conv_layer, build_norm_layer, - constant_init, kaiming_init) -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _BatchNorm - -from comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ( + ConvModule, + build_conv_layer, + build_norm_layer, + constant_init, + kaiming_init, +) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _BatchNorm + +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger from ..builder import BACKBONES @@ -31,11 +36,13 @@ def __init__(self, channel, reduction=16, with_cp=False): self.with_cp = with_cp self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( - nn.Linear(channel, channel // reduction), nn.ReLU(inplace=True), - nn.Linear(channel // reduction, channel), nn.Sigmoid()) + nn.Linear(channel, channel // reduction), + nn.ReLU(inplace=True), + nn.Linear(channel // reduction, channel), + nn.Sigmoid(), + ) def forward(self, x): - def _inner_forward(x): num_batch, num_channel = x.size()[:2] y = self.avg_pool(x).view(num_batch, num_channel) @@ -75,46 +82,37 @@ class ContextGuidedBlock(nn.Module): memory while slowing down the training speed. Default: False. """ - def __init__(self, - in_channels, - out_channels, - dilation=2, - reduction=16, - skip_connect=True, - downsample=False, - conv_cfg=None, - norm_cfg=dict(type='BN', requires_grad=True), - act_cfg=dict(type='PReLU'), - with_cp=False): + def __init__( + self, + in_channels, + out_channels, + dilation=2, + reduction=16, + skip_connect=True, + downsample=False, + conv_cfg=None, + norm_cfg=dict(type="BN", requires_grad=True), + act_cfg=dict(type="PReLU"), + with_cp=False, + ): super(ContextGuidedBlock, self).__init__() self.with_cp = with_cp self.downsample = downsample channels = out_channels if downsample else out_channels // 2 - if 'type' in act_cfg and act_cfg['type'] == 'PReLU': - act_cfg['num_parameters'] = channels + if "type" in act_cfg and act_cfg["type"] == "PReLU": + act_cfg["num_parameters"] = channels kernel_size = 3 if downsample else 1 stride = 2 if downsample else 1 padding = (kernel_size - 1) // 2 self.conv1x1 = ConvModule( - in_channels, - channels, - kernel_size, - stride, - padding, - conv_cfg=conv_cfg, - norm_cfg=norm_cfg, - act_cfg=act_cfg) + in_channels, channels, kernel_size, stride, padding, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg + ) self.f_loc = build_conv_layer( - conv_cfg, - channels, - channels, - kernel_size=3, - padding=1, - groups=channels, - bias=False) + conv_cfg, channels, channels, kernel_size=3, padding=1, groups=channels, bias=False + ) self.f_sur = build_conv_layer( conv_cfg, channels, @@ -123,24 +121,19 @@ def __init__(self, padding=dilation, groups=channels, dilation=dilation, - bias=False) + bias=False, + ) self.bn = build_norm_layer(norm_cfg, 2 * channels)[1] self.activate = nn.PReLU(2 * channels) if downsample: - self.bottleneck = build_conv_layer( - conv_cfg, - 2 * channels, - out_channels, - kernel_size=1, - bias=False) + self.bottleneck = build_conv_layer(conv_cfg, 2 * channels, out_channels, kernel_size=1, bias=False) self.skip_connect = skip_connect and not downsample self.f_glo = GlobalContextExtractor(out_channels, reduction, with_cp) def forward(self, x): - def _inner_forward(x): out = self.conv1x1(x) loc = self.f_loc(out) @@ -212,23 +205,23 @@ class CGNet(nn.Module): memory while slowing down the training speed. Default: False. """ - def __init__(self, - in_channels=3, - num_channels=(32, 64, 128), - num_blocks=(3, 21), - dilations=(2, 4), - reductions=(8, 16), - conv_cfg=None, - norm_cfg=dict(type='BN', requires_grad=True), - act_cfg=dict(type='PReLU'), - norm_eval=False, - with_cp=False): - + def __init__( + self, + in_channels=3, + num_channels=(32, 64, 128), + num_blocks=(3, 21), + dilations=(2, 4), + reductions=(8, 16), + conv_cfg=None, + norm_cfg=dict(type="BN", requires_grad=True), + act_cfg=dict(type="PReLU"), + norm_eval=False, + with_cp=False, + ): super(CGNet, self).__init__() self.in_channels = in_channels self.num_channels = num_channels - assert isinstance(self.num_channels, tuple) and len( - self.num_channels) == 3 + assert isinstance(self.num_channels, tuple) and len(self.num_channels) == 3 self.num_blocks = num_blocks assert isinstance(self.num_blocks, tuple) and len(self.num_blocks) == 2 self.dilations = dilations @@ -238,8 +231,8 @@ def __init__(self, self.conv_cfg = conv_cfg self.norm_cfg = norm_cfg self.act_cfg = act_cfg - if 'type' in self.act_cfg and self.act_cfg['type'] == 'PReLU': - self.act_cfg['num_parameters'] = num_channels[0] + if "type" in self.act_cfg and self.act_cfg["type"] == "PReLU": + self.act_cfg["num_parameters"] = num_channels[0] self.norm_eval = norm_eval self.with_cp = with_cp @@ -255,16 +248,16 @@ def __init__(self, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg)) + act_cfg=act_cfg, + ) + ) cur_channels = num_channels[0] self.inject_2x = InputInjection(1) # down-sample for Input, factor=2 self.inject_4x = InputInjection(2) # down-sample for Input, factor=4 cur_channels += in_channels - self.norm_prelu_0 = nn.Sequential( - build_norm_layer(norm_cfg, cur_channels)[1], - nn.PReLU(cur_channels)) + self.norm_prelu_0 = nn.Sequential(build_norm_layer(norm_cfg, cur_channels)[1], nn.PReLU(cur_channels)) # stage 1 self.level1 = nn.ModuleList() @@ -279,12 +272,12 @@ def __init__(self, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg, - with_cp=with_cp)) # CG block + with_cp=with_cp, + ) + ) # CG block cur_channels = 2 * num_channels[1] + in_channels - self.norm_prelu_1 = nn.Sequential( - build_norm_layer(norm_cfg, cur_channels)[1], - nn.PReLU(cur_channels)) + self.norm_prelu_1 = nn.Sequential(build_norm_layer(norm_cfg, cur_channels)[1], nn.PReLU(cur_channels)) # stage 2 self.level2 = nn.ModuleList() @@ -299,12 +292,12 @@ def __init__(self, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg, - with_cp=with_cp)) # CG block + with_cp=with_cp, + ) + ) # CG block cur_channels = 2 * num_channels[2] - self.norm_prelu_2 = nn.Sequential( - build_norm_layer(norm_cfg, cur_channels)[1], - nn.PReLU(cur_channels)) + self.norm_prelu_2 = nn.Sequential(build_norm_layer(norm_cfg, cur_channels)[1], nn.PReLU(cur_channels)) def forward(self, x): output = [] @@ -354,7 +347,7 @@ def init_weights(self, pretrained=None): elif isinstance(m, nn.PReLU): constant_init(m, 0) else: - raise TypeError('pretrained must be a str or None') + raise TypeError("pretrained must be a str or None") def train(self, mode=True): """Convert the model into training mode will keeping the normalization diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/fast_scnn.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/fast_scnn.py index 78f97e0a..bace4bfc 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/fast_scnn.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/fast_scnn.py @@ -1,11 +1,15 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import (ConvModule, DepthwiseSeparableConvModule, constant_init, - kaiming_init) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ( + ConvModule, + DepthwiseSeparableConvModule, + constant_init, + kaiming_init, +) from torch.nn.modules.batchnorm import _BatchNorm -from comfy_controlnet_preprocessors.uniformer.mmseg.models.decode_heads.psp_head import PPM -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.models.decode_heads.psp_head import PPM +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import BACKBONES from ..utils.inverted_residual import InvertedResidual @@ -26,13 +30,15 @@ class LearningToDownsample(nn.Module): dict(type='ReLU') """ - def __init__(self, - in_channels, - dw_channels, - out_channels, - conv_cfg=None, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU')): + def __init__( + self, + in_channels, + dw_channels, + out_channels, + conv_cfg=None, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU"), + ): super(LearningToDownsample, self).__init__() self.conv_cfg = conv_cfg self.norm_cfg = norm_cfg @@ -47,21 +53,14 @@ def __init__(self, stride=2, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.dsconv1 = DepthwiseSeparableConvModule( - dw_channels1, - dw_channels2, - kernel_size=3, - stride=2, - padding=1, - norm_cfg=self.norm_cfg) + dw_channels1, dw_channels2, kernel_size=3, stride=2, padding=1, norm_cfg=self.norm_cfg + ) self.dsconv2 = DepthwiseSeparableConvModule( - dw_channels2, - out_channels, - kernel_size=3, - stride=2, - padding=1, - norm_cfg=self.norm_cfg) + dw_channels2, out_channels, kernel_size=3, stride=2, padding=1, norm_cfg=self.norm_cfg + ) def forward(self, x): x = self.conv(x) @@ -103,32 +102,32 @@ class GlobalFeatureExtractor(nn.Module): Default: False """ - def __init__(self, - in_channels=64, - block_channels=(64, 96, 128), - out_channels=128, - expand_ratio=6, - num_blocks=(3, 3, 3), - strides=(2, 2, 1), - pool_scales=(1, 2, 3, 6), - conv_cfg=None, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU'), - align_corners=False): + def __init__( + self, + in_channels=64, + block_channels=(64, 96, 128), + out_channels=128, + expand_ratio=6, + num_blocks=(3, 3, 3), + strides=(2, 2, 1), + pool_scales=(1, 2, 3, 6), + conv_cfg=None, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU"), + align_corners=False, + ): super(GlobalFeatureExtractor, self).__init__() self.conv_cfg = conv_cfg self.norm_cfg = norm_cfg self.act_cfg = act_cfg assert len(block_channels) == len(num_blocks) == 3 - self.bottleneck1 = self._make_layer(in_channels, block_channels[0], - num_blocks[0], strides[0], - expand_ratio) - self.bottleneck2 = self._make_layer(block_channels[0], - block_channels[1], num_blocks[1], - strides[1], expand_ratio) - self.bottleneck3 = self._make_layer(block_channels[1], - block_channels[2], num_blocks[2], - strides[2], expand_ratio) + self.bottleneck1 = self._make_layer(in_channels, block_channels[0], num_blocks[0], strides[0], expand_ratio) + self.bottleneck2 = self._make_layer( + block_channels[0], block_channels[1], num_blocks[1], strides[1], expand_ratio + ) + self.bottleneck3 = self._make_layer( + block_channels[1], block_channels[2], num_blocks[2], strides[2], expand_ratio + ) self.ppm = PPM( pool_scales, block_channels[2], @@ -136,37 +135,21 @@ def __init__(self, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg, - align_corners=align_corners) + align_corners=align_corners, + ) self.out = ConvModule( block_channels[2] * 2, out_channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) - def _make_layer(self, - in_channels, - out_channels, - blocks, - stride=1, - expand_ratio=6): - layers = [ - InvertedResidual( - in_channels, - out_channels, - stride, - expand_ratio, - norm_cfg=self.norm_cfg) - ] + def _make_layer(self, in_channels, out_channels, blocks, stride=1, expand_ratio=6): + layers = [InvertedResidual(in_channels, out_channels, stride, expand_ratio, norm_cfg=self.norm_cfg)] for i in range(1, blocks): - layers.append( - InvertedResidual( - out_channels, - out_channels, - 1, - expand_ratio, - norm_cfg=self.norm_cfg)) + layers.append(InvertedResidual(out_channels, out_channels, 1, expand_ratio, norm_cfg=self.norm_cfg)) return nn.Sequential(*layers) def forward(self, x): @@ -196,48 +179,36 @@ class FeatureFusionModule(nn.Module): Default: False """ - def __init__(self, - higher_in_channels, - lower_in_channels, - out_channels, - conv_cfg=None, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU'), - align_corners=False): + def __init__( + self, + higher_in_channels, + lower_in_channels, + out_channels, + conv_cfg=None, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU"), + align_corners=False, + ): super(FeatureFusionModule, self).__init__() self.conv_cfg = conv_cfg self.norm_cfg = norm_cfg self.act_cfg = act_cfg self.align_corners = align_corners self.dwconv = ConvModule( - lower_in_channels, - out_channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + lower_in_channels, out_channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) self.conv_lower_res = ConvModule( - out_channels, - out_channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=None) + out_channels, out_channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=None + ) self.conv_higher_res = ConvModule( - higher_in_channels, - out_channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=None) + higher_in_channels, out_channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=None + ) self.relu = nn.ReLU(True) def forward(self, higher_res_feature, lower_res_feature): lower_res_feature = resize( - lower_res_feature, - size=higher_res_feature.size()[2:], - mode='bilinear', - align_corners=self.align_corners) + lower_res_feature, size=higher_res_feature.size()[2:], mode="bilinear", align_corners=self.align_corners + ) lower_res_feature = self.dwconv(lower_res_feature) lower_res_feature = self.conv_lower_res(lower_res_feature) @@ -293,29 +264,34 @@ class FastSCNN(nn.Module): Default: False """ - def __init__(self, - in_channels=3, - downsample_dw_channels=(32, 48), - global_in_channels=64, - global_block_channels=(64, 96, 128), - global_block_strides=(2, 2, 1), - global_out_channels=128, - higher_in_channels=64, - lower_in_channels=128, - fusion_out_channels=128, - out_indices=(0, 1, 2), - conv_cfg=None, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU'), - align_corners=False): - + def __init__( + self, + in_channels=3, + downsample_dw_channels=(32, 48), + global_in_channels=64, + global_block_channels=(64, 96, 128), + global_block_strides=(2, 2, 1), + global_out_channels=128, + higher_in_channels=64, + lower_in_channels=128, + fusion_out_channels=128, + out_indices=(0, 1, 2), + conv_cfg=None, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU"), + align_corners=False, + ): super(FastSCNN, self).__init__() if global_in_channels != higher_in_channels: - raise AssertionError('Global Input Channels must be the same \ - with Higher Input Channels!') + raise AssertionError( + "Global Input Channels must be the same \ + with Higher Input Channels!" + ) elif global_out_channels != lower_in_channels: - raise AssertionError('Global Output Channels must be the same \ - with Lower Input Channels!') + raise AssertionError( + "Global Output Channels must be the same \ + with Lower Input Channels!" + ) self.in_channels = in_channels self.downsample_dw_channels1 = downsample_dw_channels[0] @@ -338,7 +314,8 @@ def __init__(self, global_in_channels, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.global_feature_extractor = GlobalFeatureExtractor( global_in_channels, global_block_channels, @@ -347,7 +324,8 @@ def __init__(self, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg, - align_corners=self.align_corners) + align_corners=self.align_corners, + ) self.feature_fusion = FeatureFusionModule( higher_in_channels, lower_in_channels, @@ -355,7 +333,8 @@ def __init__(self, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg, - align_corners=self.align_corners) + align_corners=self.align_corners, + ) def init_weights(self, pretrained=None): for m in self.modules(): @@ -367,8 +346,7 @@ def init_weights(self, pretrained=None): def forward(self, x): higher_res_features = self.learning_to_downsample(x) lower_res_features = self.global_feature_extractor(higher_res_features) - fusion_output = self.feature_fusion(higher_res_features, - lower_res_features) + fusion_output = self.feature_fusion(higher_res_features, lower_res_features) outs = [higher_res_features, lower_res_features, fusion_output] outs = [outs[i] for i in self.out_indices] diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/hrnet.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/hrnet.py index d20a9c11..fbda0e02 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/hrnet.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/hrnet.py @@ -1,11 +1,15 @@ import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import (build_conv_layer, build_norm_layer, constant_init, - kaiming_init) -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _BatchNorm - -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import Upsample, resize -from comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ( + build_conv_layer, + build_norm_layer, + constant_init, + kaiming_init, +) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _BatchNorm + +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import Upsample, resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger from ..builder import BACKBONES from .resnet import BasicBlock, Bottleneck @@ -17,19 +21,20 @@ class HRModule(nn.Module): is in this module. """ - def __init__(self, - num_branches, - blocks, - num_blocks, - in_channels, - num_channels, - multiscale_output=True, - with_cp=False, - conv_cfg=None, - norm_cfg=dict(type='BN', requires_grad=True)): + def __init__( + self, + num_branches, + blocks, + num_blocks, + in_channels, + num_channels, + multiscale_output=True, + with_cp=False, + conv_cfg=None, + norm_cfg=dict(type="BN", requires_grad=True), + ): super(HRModule, self).__init__() - self._check_branches(num_branches, num_blocks, in_channels, - num_channels) + self._check_branches(num_branches, num_blocks, in_channels, num_channels) self.in_channels = in_channels self.num_branches = num_branches @@ -38,40 +43,28 @@ def __init__(self, self.norm_cfg = norm_cfg self.conv_cfg = conv_cfg self.with_cp = with_cp - self.branches = self._make_branches(num_branches, blocks, num_blocks, - num_channels) + self.branches = self._make_branches(num_branches, blocks, num_blocks, num_channels) self.fuse_layers = self._make_fuse_layers() self.relu = nn.ReLU(inplace=False) - def _check_branches(self, num_branches, num_blocks, in_channels, - num_channels): + def _check_branches(self, num_branches, num_blocks, in_channels, num_channels): """Check branches configuration.""" if num_branches != len(num_blocks): - error_msg = f'NUM_BRANCHES({num_branches}) <> NUM_BLOCKS(' \ - f'{len(num_blocks)})' + error_msg = f"NUM_BRANCHES({num_branches}) <> NUM_BLOCKS(" f"{len(num_blocks)})" raise ValueError(error_msg) if num_branches != len(num_channels): - error_msg = f'NUM_BRANCHES({num_branches}) <> NUM_CHANNELS(' \ - f'{len(num_channels)})' + error_msg = f"NUM_BRANCHES({num_branches}) <> NUM_CHANNELS(" f"{len(num_channels)})" raise ValueError(error_msg) if num_branches != len(in_channels): - error_msg = f'NUM_BRANCHES({num_branches}) <> NUM_INCHANNELS(' \ - f'{len(in_channels)})' + error_msg = f"NUM_BRANCHES({num_branches}) <> NUM_INCHANNELS(" f"{len(in_channels)})" raise ValueError(error_msg) - def _make_one_branch(self, - branch_index, - block, - num_blocks, - num_channels, - stride=1): + def _make_one_branch(self, branch_index, block, num_blocks, num_channels, stride=1): """Build one branch.""" downsample = None - if stride != 1 or \ - self.in_channels[branch_index] != \ - num_channels[branch_index] * block.expansion: + if stride != 1 or self.in_channels[branch_index] != num_channels[branch_index] * block.expansion: downsample = nn.Sequential( build_conv_layer( self.conv_cfg, @@ -79,9 +72,10 @@ def _make_one_branch(self, num_channels[branch_index] * block.expansion, kernel_size=1, stride=stride, - bias=False), - build_norm_layer(self.norm_cfg, num_channels[branch_index] * - block.expansion)[1]) + bias=False, + ), + build_norm_layer(self.norm_cfg, num_channels[branch_index] * block.expansion)[1], + ) layers = [] layers.append( @@ -92,9 +86,10 @@ def _make_one_branch(self, downsample=downsample, with_cp=self.with_cp, norm_cfg=self.norm_cfg, - conv_cfg=self.conv_cfg)) - self.in_channels[branch_index] = \ - num_channels[branch_index] * block.expansion + conv_cfg=self.conv_cfg, + ) + ) + self.in_channels[branch_index] = num_channels[branch_index] * block.expansion for i in range(1, num_blocks[branch_index]): layers.append( block( @@ -102,7 +97,9 @@ def _make_one_branch(self, num_channels[branch_index], with_cp=self.with_cp, norm_cfg=self.norm_cfg, - conv_cfg=self.conv_cfg)) + conv_cfg=self.conv_cfg, + ) + ) return nn.Sequential(*layers) @@ -111,8 +108,7 @@ def _make_branches(self, num_branches, block, num_blocks, num_channels): branches = [] for i in range(num_branches): - branches.append( - self._make_one_branch(i, block, num_blocks, num_channels)) + branches.append(self._make_one_branch(i, block, num_blocks, num_channels)) return nn.ModuleList(branches) @@ -138,13 +134,13 @@ def _make_fuse_layers(self): kernel_size=1, stride=1, padding=0, - bias=False), + bias=False, + ), build_norm_layer(self.norm_cfg, in_channels[i])[1], # we set align_corners=False for HRNet - Upsample( - scale_factor=2**(j - i), - mode='bilinear', - align_corners=False))) + Upsample(scale_factor=2 ** (j - i), mode="bilinear", align_corners=False), + ) + ) elif j == i: fuse_layer.append(None) else: @@ -160,9 +156,11 @@ def _make_fuse_layers(self): kernel_size=3, stride=2, padding=1, - bias=False), - build_norm_layer(self.norm_cfg, - in_channels[i])[1])) + bias=False, + ), + build_norm_layer(self.norm_cfg, in_channels[i])[1], + ) + ) else: conv_downsamples.append( nn.Sequential( @@ -173,10 +171,12 @@ def _make_fuse_layers(self): kernel_size=3, stride=2, padding=1, - bias=False), - build_norm_layer(self.norm_cfg, - in_channels[j])[1], - nn.ReLU(inplace=False))) + bias=False, + ), + build_norm_layer(self.norm_cfg, in_channels[j])[1], + nn.ReLU(inplace=False), + ) + ) fuse_layer.append(nn.Sequential(*conv_downsamples)) fuse_layers.append(nn.ModuleList(fuse_layer)) @@ -198,10 +198,8 @@ def forward(self, x): y += x[j] elif j > i: y = y + resize( - self.fuse_layers[i][j](x[j]), - size=x[i].shape[2:], - mode='bilinear', - align_corners=False) + self.fuse_layers[i][j](x[j]), size=x[i].shape[2:], mode="bilinear", align_corners=False + ) else: y += self.fuse_layers[i][j](x[j]) x_fuse.append(self.relu(y)) @@ -229,7 +227,7 @@ class HRNet(nn.Module): in resblocks to let them behave as identity. Example: - >>> from comfy_controlnet_preprocessors.uniformer.mmseg.models import HRNet + >>> from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.models import HRNet >>> import torch >>> extra = dict( >>> stage1=dict( @@ -268,16 +266,18 @@ class HRNet(nn.Module): (1, 256, 1, 1) """ - blocks_dict = {'BASIC': BasicBlock, 'BOTTLENECK': Bottleneck} - - def __init__(self, - extra, - in_channels=3, - conv_cfg=None, - norm_cfg=dict(type='BN', requires_grad=True), - norm_eval=False, - with_cp=False, - zero_init_residual=False): + blocks_dict = {"BASIC": BasicBlock, "BOTTLENECK": Bottleneck} + + def __init__( + self, + extra, + in_channels=3, + conv_cfg=None, + norm_cfg=dict(type="BN", requires_grad=True), + norm_eval=False, + with_cp=False, + zero_init_residual=False, + ): super(HRNet, self).__init__() self.extra = extra self.conv_cfg = conv_cfg @@ -290,73 +290,53 @@ def __init__(self, self.norm1_name, norm1 = build_norm_layer(self.norm_cfg, 64, postfix=1) self.norm2_name, norm2 = build_norm_layer(self.norm_cfg, 64, postfix=2) - self.conv1 = build_conv_layer( - self.conv_cfg, - in_channels, - 64, - kernel_size=3, - stride=2, - padding=1, - bias=False) + self.conv1 = build_conv_layer(self.conv_cfg, in_channels, 64, kernel_size=3, stride=2, padding=1, bias=False) self.add_module(self.norm1_name, norm1) - self.conv2 = build_conv_layer( - self.conv_cfg, - 64, - 64, - kernel_size=3, - stride=2, - padding=1, - bias=False) + self.conv2 = build_conv_layer(self.conv_cfg, 64, 64, kernel_size=3, stride=2, padding=1, bias=False) self.add_module(self.norm2_name, norm2) self.relu = nn.ReLU(inplace=True) # stage 1 - self.stage1_cfg = self.extra['stage1'] - num_channels = self.stage1_cfg['num_channels'][0] - block_type = self.stage1_cfg['block'] - num_blocks = self.stage1_cfg['num_blocks'][0] + self.stage1_cfg = self.extra["stage1"] + num_channels = self.stage1_cfg["num_channels"][0] + block_type = self.stage1_cfg["block"] + num_blocks = self.stage1_cfg["num_blocks"][0] block = self.blocks_dict[block_type] stage1_out_channels = num_channels * block.expansion self.layer1 = self._make_layer(block, 64, num_channels, num_blocks) # stage 2 - self.stage2_cfg = self.extra['stage2'] - num_channels = self.stage2_cfg['num_channels'] - block_type = self.stage2_cfg['block'] + self.stage2_cfg = self.extra["stage2"] + num_channels = self.stage2_cfg["num_channels"] + block_type = self.stage2_cfg["block"] block = self.blocks_dict[block_type] num_channels = [channel * block.expansion for channel in num_channels] - self.transition1 = self._make_transition_layer([stage1_out_channels], - num_channels) - self.stage2, pre_stage_channels = self._make_stage( - self.stage2_cfg, num_channels) + self.transition1 = self._make_transition_layer([stage1_out_channels], num_channels) + self.stage2, pre_stage_channels = self._make_stage(self.stage2_cfg, num_channels) # stage 3 - self.stage3_cfg = self.extra['stage3'] - num_channels = self.stage3_cfg['num_channels'] - block_type = self.stage3_cfg['block'] + self.stage3_cfg = self.extra["stage3"] + num_channels = self.stage3_cfg["num_channels"] + block_type = self.stage3_cfg["block"] block = self.blocks_dict[block_type] num_channels = [channel * block.expansion for channel in num_channels] - self.transition2 = self._make_transition_layer(pre_stage_channels, - num_channels) - self.stage3, pre_stage_channels = self._make_stage( - self.stage3_cfg, num_channels) + self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels) + self.stage3, pre_stage_channels = self._make_stage(self.stage3_cfg, num_channels) # stage 4 - self.stage4_cfg = self.extra['stage4'] - num_channels = self.stage4_cfg['num_channels'] - block_type = self.stage4_cfg['block'] + self.stage4_cfg = self.extra["stage4"] + num_channels = self.stage4_cfg["num_channels"] + block_type = self.stage4_cfg["block"] block = self.blocks_dict[block_type] num_channels = [channel * block.expansion for channel in num_channels] - self.transition3 = self._make_transition_layer(pre_stage_channels, - num_channels) - self.stage4, pre_stage_channels = self._make_stage( - self.stage4_cfg, num_channels) + self.transition3 = self._make_transition_layer(pre_stage_channels, num_channels) + self.stage4, pre_stage_channels = self._make_stage(self.stage4_cfg, num_channels) @property def norm1(self): @@ -368,8 +348,7 @@ def norm2(self): """nn.Module: the normalization layer named "norm2" """ return getattr(self, self.norm2_name) - def _make_transition_layer(self, num_channels_pre_layer, - num_channels_cur_layer): + def _make_transition_layer(self, num_channels_pre_layer, num_channels_cur_layer): """Make transition layer.""" num_branches_cur = len(num_channels_cur_layer) num_branches_pre = len(num_channels_pre_layer) @@ -387,18 +366,19 @@ def _make_transition_layer(self, num_channels_pre_layer, kernel_size=3, stride=1, padding=1, - bias=False), - build_norm_layer(self.norm_cfg, - num_channels_cur_layer[i])[1], - nn.ReLU(inplace=True))) + bias=False, + ), + build_norm_layer(self.norm_cfg, num_channels_cur_layer[i])[1], + nn.ReLU(inplace=True), + ) + ) else: transition_layers.append(None) else: conv_downsamples = [] for j in range(i + 1 - num_branches_pre): in_channels = num_channels_pre_layer[-1] - out_channels = num_channels_cur_layer[i] \ - if j == i - num_branches_pre else in_channels + out_channels = num_channels_cur_layer[i] if j == i - num_branches_pre else in_channels conv_downsamples.append( nn.Sequential( build_conv_layer( @@ -408,9 +388,12 @@ def _make_transition_layer(self, num_channels_pre_layer, kernel_size=3, stride=2, padding=1, - bias=False), + bias=False, + ), build_norm_layer(self.norm_cfg, out_channels)[1], - nn.ReLU(inplace=True))) + nn.ReLU(inplace=True), + ) + ) transition_layers.append(nn.Sequential(*conv_downsamples)) return nn.ModuleList(transition_layers) @@ -421,13 +404,10 @@ def _make_layer(self, block, inplanes, planes, blocks, stride=1): if stride != 1 or inplanes != planes * block.expansion: downsample = nn.Sequential( build_conv_layer( - self.conv_cfg, - inplanes, - planes * block.expansion, - kernel_size=1, - stride=stride, - bias=False), - build_norm_layer(self.norm_cfg, planes * block.expansion)[1]) + self.conv_cfg, inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False + ), + build_norm_layer(self.norm_cfg, planes * block.expansion)[1], + ) layers = [] layers.append( @@ -438,26 +418,24 @@ def _make_layer(self, block, inplanes, planes, blocks, stride=1): downsample=downsample, with_cp=self.with_cp, norm_cfg=self.norm_cfg, - conv_cfg=self.conv_cfg)) + conv_cfg=self.conv_cfg, + ) + ) inplanes = planes * block.expansion for i in range(1, blocks): layers.append( - block( - inplanes, - planes, - with_cp=self.with_cp, - norm_cfg=self.norm_cfg, - conv_cfg=self.conv_cfg)) + block(inplanes, planes, with_cp=self.with_cp, norm_cfg=self.norm_cfg, conv_cfg=self.conv_cfg) + ) return nn.Sequential(*layers) def _make_stage(self, layer_config, in_channels, multiscale_output=True): """Make each stage.""" - num_modules = layer_config['num_modules'] - num_branches = layer_config['num_branches'] - num_blocks = layer_config['num_blocks'] - num_channels = layer_config['num_channels'] - block = self.blocks_dict[layer_config['block']] + num_modules = layer_config["num_modules"] + num_branches = layer_config["num_branches"] + num_blocks = layer_config["num_blocks"] + num_channels = layer_config["num_channels"] + block = self.blocks_dict[layer_config["block"]] hr_modules = [] for i in range(num_modules): @@ -477,7 +455,9 @@ def _make_stage(self, layer_config, in_channels, multiscale_output=True): reset_multiscale_output, with_cp=self.with_cp, norm_cfg=self.norm_cfg, - conv_cfg=self.conv_cfg)) + conv_cfg=self.conv_cfg, + ) + ) return nn.Sequential(*hr_modules), in_channels @@ -505,7 +485,7 @@ def init_weights(self, pretrained=None): elif isinstance(m, BasicBlock): constant_init(m.norm2, 0) else: - raise TypeError('pretrained must be a str or None') + raise TypeError("pretrained must be a str or None") def forward(self, x): """Forward function.""" @@ -519,7 +499,7 @@ def forward(self, x): x = self.layer1(x) x_list = [] - for i in range(self.stage2_cfg['num_branches']): + for i in range(self.stage2_cfg["num_branches"]): if self.transition1[i] is not None: x_list.append(self.transition1[i](x)) else: @@ -527,7 +507,7 @@ def forward(self, x): y_list = self.stage2(x_list) x_list = [] - for i in range(self.stage3_cfg['num_branches']): + for i in range(self.stage3_cfg["num_branches"]): if self.transition2[i] is not None: x_list.append(self.transition2[i](y_list[-1])) else: @@ -535,7 +515,7 @@ def forward(self, x): y_list = self.stage3(x_list) x_list = [] - for i in range(self.stage4_cfg['num_branches']): + for i in range(self.stage4_cfg["num_branches"]): if self.transition3[i] is not None: x_list.append(self.transition3[i](y_list[-1])) else: diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/mobilenet_v2.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/mobilenet_v2.py index 080cce42..b286c9e1 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/mobilenet_v2.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/mobilenet_v2.py @@ -1,8 +1,8 @@ import logging import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, constant_init, kaiming_init -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, constant_init, kaiming_init +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint from torch.nn.modules.batchnorm import _BatchNorm from ..builder import BACKBONES @@ -39,20 +39,21 @@ class MobileNetV2(nn.Module): # Parameters to build layers. 3 parameters are needed to construct a # layer, from left to right: expand_ratio, channel, num_blocks. - arch_settings = [[1, 16, 1], [6, 24, 2], [6, 32, 3], [6, 64, 4], - [6, 96, 3], [6, 160, 3], [6, 320, 1]] - - def __init__(self, - widen_factor=1., - strides=(1, 2, 2, 2, 1, 2, 1), - dilations=(1, 1, 1, 1, 1, 1, 1), - out_indices=(1, 2, 4, 6), - frozen_stages=-1, - conv_cfg=None, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU6'), - norm_eval=False, - with_cp=False): + arch_settings = [[1, 16, 1], [6, 24, 2], [6, 32, 3], [6, 64, 4], [6, 96, 3], [6, 160, 3], [6, 320, 1]] + + def __init__( + self, + widen_factor=1.0, + strides=(1, 2, 2, 2, 1, 2, 1), + dilations=(1, 1, 1, 1, 1, 1, 1), + out_indices=(1, 2, 4, 6), + frozen_stages=-1, + conv_cfg=None, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU6"), + norm_eval=False, + with_cp=False, + ): super(MobileNetV2, self).__init__() self.widen_factor = widen_factor self.strides = strides @@ -61,12 +62,10 @@ def __init__(self, self.out_indices = out_indices for index in out_indices: if index not in range(0, 7): - raise ValueError('the item in out_indices must in ' - f'range(0, 8). But received {index}') + raise ValueError("the item in out_indices must in " f"range(0, 8). But received {index}") if frozen_stages not in range(-1, 7): - raise ValueError('frozen_stages must be in range(-1, 7). ' - f'But received {frozen_stages}') + raise ValueError("frozen_stages must be in range(-1, 7). " f"But received {frozen_stages}") self.out_indices = out_indices self.frozen_stages = frozen_stages self.conv_cfg = conv_cfg @@ -85,7 +84,8 @@ def __init__(self, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.layers = [] @@ -99,13 +99,13 @@ def __init__(self, num_blocks=num_blocks, stride=stride, dilation=dilation, - expand_ratio=expand_ratio) - layer_name = f'layer{i + 1}' + expand_ratio=expand_ratio, + ) + layer_name = f"layer{i + 1}" self.add_module(layer_name, inverted_res_layer) self.layers.append(layer_name) - def make_layer(self, out_channels, num_blocks, stride, dilation, - expand_ratio): + def make_layer(self, out_channels, num_blocks, stride, dilation, expand_ratio): """Stack InvertedResidual blocks to build a layer for MobileNetV2. Args: @@ -128,7 +128,9 @@ def make_layer(self, out_channels, num_blocks, stride, dilation, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg, - with_cp=self.with_cp)) + with_cp=self.with_cp, + ) + ) self.in_channels = out_channels return nn.Sequential(*layers) @@ -144,7 +146,7 @@ def init_weights(self, pretrained=None): elif isinstance(m, (_BatchNorm, nn.GroupNorm)): constant_init(m, 1) else: - raise TypeError('pretrained must be a str or None') + raise TypeError("pretrained must be a str or None") def forward(self, x): x = self.conv1(x) @@ -166,7 +168,7 @@ def _freeze_stages(self): for param in self.conv1.parameters(): param.requires_grad = False for i in range(1, self.frozen_stages + 1): - layer = getattr(self, f'layer{i}') + layer = getattr(self, f"layer{i}") layer.eval() for param in layer.parameters(): param.requires_grad = False diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/mobilenet_v3.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/mobilenet_v3.py index 3ef21dbb..137363e3 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/mobilenet_v3.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/mobilenet_v3.py @@ -1,10 +1,10 @@ import logging -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, constant_init, kaiming_init -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn.bricks import Conv2dAdaptivePadding -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, constant_init, kaiming_init +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn.bricks import Conv2dAdaptivePadding +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint from torch.nn.modules.batchnorm import _BatchNorm from ..builder import BACKBONES @@ -36,46 +36,53 @@ class MobileNetV3(nn.Module): some memory while slowing down the training speed. Default: False. """ + # Parameters to build each block: # [kernel size, mid channels, out channels, with_se, act type, stride] arch_settings = { - 'small': [[3, 16, 16, True, 'ReLU', 2], # block0 layer1 os=4 - [3, 72, 24, False, 'ReLU', 2], # block1 layer2 os=8 - [3, 88, 24, False, 'ReLU', 1], - [5, 96, 40, True, 'HSwish', 2], # block2 layer4 os=16 - [5, 240, 40, True, 'HSwish', 1], - [5, 240, 40, True, 'HSwish', 1], - [5, 120, 48, True, 'HSwish', 1], # block3 layer7 os=16 - [5, 144, 48, True, 'HSwish', 1], - [5, 288, 96, True, 'HSwish', 2], # block4 layer9 os=32 - [5, 576, 96, True, 'HSwish', 1], - [5, 576, 96, True, 'HSwish', 1]], - 'large': [[3, 16, 16, False, 'ReLU', 1], # block0 layer1 os=2 - [3, 64, 24, False, 'ReLU', 2], # block1 layer2 os=4 - [3, 72, 24, False, 'ReLU', 1], - [5, 72, 40, True, 'ReLU', 2], # block2 layer4 os=8 - [5, 120, 40, True, 'ReLU', 1], - [5, 120, 40, True, 'ReLU', 1], - [3, 240, 80, False, 'HSwish', 2], # block3 layer7 os=16 - [3, 200, 80, False, 'HSwish', 1], - [3, 184, 80, False, 'HSwish', 1], - [3, 184, 80, False, 'HSwish', 1], - [3, 480, 112, True, 'HSwish', 1], # block4 layer11 os=16 - [3, 672, 112, True, 'HSwish', 1], - [5, 672, 160, True, 'HSwish', 2], # block5 layer13 os=32 - [5, 960, 160, True, 'HSwish', 1], - [5, 960, 160, True, 'HSwish', 1]] + "small": [ + [3, 16, 16, True, "ReLU", 2], # block0 layer1 os=4 + [3, 72, 24, False, "ReLU", 2], # block1 layer2 os=8 + [3, 88, 24, False, "ReLU", 1], + [5, 96, 40, True, "HSwish", 2], # block2 layer4 os=16 + [5, 240, 40, True, "HSwish", 1], + [5, 240, 40, True, "HSwish", 1], + [5, 120, 48, True, "HSwish", 1], # block3 layer7 os=16 + [5, 144, 48, True, "HSwish", 1], + [5, 288, 96, True, "HSwish", 2], # block4 layer9 os=32 + [5, 576, 96, True, "HSwish", 1], + [5, 576, 96, True, "HSwish", 1], + ], + "large": [ + [3, 16, 16, False, "ReLU", 1], # block0 layer1 os=2 + [3, 64, 24, False, "ReLU", 2], # block1 layer2 os=4 + [3, 72, 24, False, "ReLU", 1], + [5, 72, 40, True, "ReLU", 2], # block2 layer4 os=8 + [5, 120, 40, True, "ReLU", 1], + [5, 120, 40, True, "ReLU", 1], + [3, 240, 80, False, "HSwish", 2], # block3 layer7 os=16 + [3, 200, 80, False, "HSwish", 1], + [3, 184, 80, False, "HSwish", 1], + [3, 184, 80, False, "HSwish", 1], + [3, 480, 112, True, "HSwish", 1], # block4 layer11 os=16 + [3, 672, 112, True, "HSwish", 1], + [5, 672, 160, True, "HSwish", 2], # block5 layer13 os=32 + [5, 960, 160, True, "HSwish", 1], + [5, 960, 160, True, "HSwish", 1], + ], } # yapf: disable - def __init__(self, - arch='small', - conv_cfg=None, - norm_cfg=dict(type='BN'), - out_indices=(0, 1, 12), - frozen_stages=-1, - reduction_factor=1, - norm_eval=False, - with_cp=False): + def __init__( + self, + arch="small", + conv_cfg=None, + norm_cfg=dict(type="BN"), + out_indices=(0, 1, 12), + frozen_stages=-1, + reduction_factor=1, + norm_eval=False, + with_cp=False, + ): super(MobileNetV3, self).__init__() assert arch in self.arch_settings assert isinstance(reduction_factor, int) and reduction_factor > 0 @@ -83,14 +90,17 @@ def __init__(self, for index in out_indices: if index not in range(0, len(self.arch_settings[arch]) + 2): raise ValueError( - 'the item in out_indices must in ' - f'range(0, {len(self.arch_settings[arch])+2}). ' - f'But received {index}') + "the item in out_indices must in " + f"range(0, {len(self.arch_settings[arch])+2}). " + f"But received {index}" + ) if frozen_stages not in range(-1, len(self.arch_settings[arch]) + 2): - raise ValueError('frozen_stages must be in range(-1, ' - f'{len(self.arch_settings[arch])+2}). ' - f'But received {frozen_stages}') + raise ValueError( + "frozen_stages must be in range(-1, " + f"{len(self.arch_settings[arch])+2}). " + f"But received {frozen_stages}" + ) self.arch = arch self.conv_cfg = conv_cfg self.norm_cfg = norm_cfg @@ -112,19 +122,18 @@ def _make_layer(self): kernel_size=3, stride=2, padding=1, - conv_cfg=dict(type='Conv2dAdaptivePadding'), + conv_cfg=dict(type="Conv2dAdaptivePadding"), norm_cfg=self.norm_cfg, - act_cfg=dict(type='HSwish')) - self.add_module('layer0', layer) - layers.append('layer0') + act_cfg=dict(type="HSwish"), + ) + self.add_module("layer0", layer) + layers.append("layer0") layer_setting = self.arch_settings[self.arch] for i, params in enumerate(layer_setting): - (kernel_size, mid_channels, out_channels, with_se, act, - stride) = params + (kernel_size, mid_channels, out_channels, with_se, act, stride) = params - if self.arch == 'large' and i >= 12 or self.arch == 'small' and \ - i >= 8: + if self.arch == "large" and i >= 12 or self.arch == "small" and i >= 8: mid_channels = mid_channels // self.reduction_factor out_channels = out_channels // self.reduction_factor @@ -132,8 +141,8 @@ def _make_layer(self): se_cfg = dict( channels=mid_channels, ratio=4, - act_cfg=(dict(type='ReLU'), - dict(type='HSigmoid', bias=3.0, divisor=6.0))) + act_cfg=(dict(type="ReLU"), dict(type="HSigmoid", bias=3.0, divisor=6.0)), + ) else: se_cfg = None @@ -148,9 +157,10 @@ def _make_layer(self): conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=dict(type=act), - with_cp=self.with_cp) + with_cp=self.with_cp, + ) in_channels = out_channels - layer_name = 'layer{}'.format(i + 1) + layer_name = "layer{}".format(i + 1) self.add_module(layer_name, layer) layers.append(layer_name) @@ -159,20 +169,21 @@ def _make_layer(self): # block6 layer16 os=32 for large model layer = ConvModule( in_channels=in_channels, - out_channels=576 if self.arch == 'small' else 960, + out_channels=576 if self.arch == "small" else 960, kernel_size=1, stride=1, dilation=4, padding=0, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=dict(type='HSwish')) - layer_name = 'layer{}'.format(len(layer_setting) + 1) + act_cfg=dict(type="HSwish"), + ) + layer_name = "layer{}".format(len(layer_setting) + 1) self.add_module(layer_name, layer) layers.append(layer_name) # next, convert backbone MobileNetV3 to a semantic segmentation version - if self.arch == 'small': + if self.arch == "small": self.layer4.depthwise_conv.conv.stride = (1, 1) self.layer9.depthwise_conv.conv.stride = (1, 1) for i in range(4, len(layers)): @@ -228,7 +239,7 @@ def init_weights(self, pretrained=None): elif isinstance(m, nn.BatchNorm2d): constant_init(m, 1) else: - raise TypeError('pretrained must be a str or None') + raise TypeError("pretrained must be a str or None") def forward(self, x): outs = [] @@ -241,7 +252,7 @@ def forward(self, x): def _freeze_stages(self): for i in range(self.frozen_stages + 1): - layer = getattr(self, f'layer{i}') + layer = getattr(self, f"layer{i}") layer.eval() for param in layer.parameters(): param.requires_grad = False diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnest.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnest.py index 06b1b601..997ebba8 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnest.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnest.py @@ -4,7 +4,7 @@ import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint as cp -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import build_conv_layer, build_norm_layer +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import build_conv_layer, build_norm_layer from ..builder import BACKBONES from ..utils import ResLayer @@ -55,19 +55,21 @@ class SplitAttentionConv2d(nn.Module): dcn (dict): Config dict for DCN. Default: None. """ - def __init__(self, - in_channels, - channels, - kernel_size, - stride=1, - padding=0, - dilation=1, - groups=1, - radix=2, - reduction_factor=4, - conv_cfg=None, - norm_cfg=dict(type='BN'), - dcn=None): + def __init__( + self, + in_channels, + channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + groups=1, + radix=2, + reduction_factor=4, + conv_cfg=None, + norm_cfg=dict(type="BN"), + dcn=None, + ): super(SplitAttentionConv2d, self).__init__() inter_channels = max(in_channels * radix // reduction_factor, 32) self.radix = radix @@ -77,9 +79,9 @@ def __init__(self, self.dcn = dcn fallback_on_stride = False if self.with_dcn: - fallback_on_stride = self.dcn.pop('fallback_on_stride', False) + fallback_on_stride = self.dcn.pop("fallback_on_stride", False) if self.with_dcn and not fallback_on_stride: - assert conv_cfg is None, 'conv_cfg must be None for DCN' + assert conv_cfg is None, "conv_cfg must be None for DCN" conv_cfg = dcn self.conv = build_conv_layer( conv_cfg, @@ -90,18 +92,15 @@ def __init__(self, padding=padding, dilation=dilation, groups=groups * radix, - bias=False) - self.norm0_name, norm0 = build_norm_layer( - norm_cfg, channels * radix, postfix=0) + bias=False, + ) + self.norm0_name, norm0 = build_norm_layer(norm_cfg, channels * radix, postfix=0) self.add_module(self.norm0_name, norm0) self.relu = nn.ReLU(inplace=True) - self.fc1 = build_conv_layer( - None, channels, inter_channels, 1, groups=self.groups) - self.norm1_name, norm1 = build_norm_layer( - norm_cfg, inter_channels, postfix=1) + self.fc1 = build_conv_layer(None, channels, inter_channels, 1, groups=self.groups) + self.norm1_name, norm1 = build_norm_layer(norm_cfg, inter_channels, postfix=1) self.add_module(self.norm1_name, norm1) - self.fc2 = build_conv_layer( - None, inter_channels, channels * radix, 1, groups=self.groups) + self.fc2 = build_conv_layer(None, inter_channels, channels * radix, 1, groups=self.groups) self.rsoftmax = RSoftmax(radix, groups) @property @@ -160,41 +159,37 @@ class Bottleneck(_Bottleneck): Bottleneck. Default: True. kwargs (dict): Key word arguments for base class. """ + expansion = 4 - def __init__(self, - inplanes, - planes, - groups=1, - base_width=4, - base_channels=64, - radix=2, - reduction_factor=4, - avg_down_stride=True, - **kwargs): + def __init__( + self, + inplanes, + planes, + groups=1, + base_width=4, + base_channels=64, + radix=2, + reduction_factor=4, + avg_down_stride=True, + **kwargs, + ): """Bottleneck block for ResNeSt.""" super(Bottleneck, self).__init__(inplanes, planes, **kwargs) if groups == 1: width = self.planes else: - width = math.floor(self.planes * - (base_width / base_channels)) * groups + width = math.floor(self.planes * (base_width / base_channels)) * groups self.avg_down_stride = avg_down_stride and self.conv2_stride > 1 - self.norm1_name, norm1 = build_norm_layer( - self.norm_cfg, width, postfix=1) - self.norm3_name, norm3 = build_norm_layer( - self.norm_cfg, self.planes * self.expansion, postfix=3) + self.norm1_name, norm1 = build_norm_layer(self.norm_cfg, width, postfix=1) + self.norm3_name, norm3 = build_norm_layer(self.norm_cfg, self.planes * self.expansion, postfix=3) self.conv1 = build_conv_layer( - self.conv_cfg, - self.inplanes, - width, - kernel_size=1, - stride=self.conv1_stride, - bias=False) + self.conv_cfg, self.inplanes, width, kernel_size=1, stride=self.conv1_stride, bias=False + ) self.add_module(self.norm1_name, norm1) self.with_modulated_dcn = False self.conv2 = SplitAttentionConv2d( @@ -209,22 +204,17 @@ def __init__(self, reduction_factor=reduction_factor, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - dcn=self.dcn) + dcn=self.dcn, + ) delattr(self, self.norm2_name) if self.avg_down_stride: self.avd_layer = nn.AvgPool2d(3, self.conv2_stride, padding=1) - self.conv3 = build_conv_layer( - self.conv_cfg, - width, - self.planes * self.expansion, - kernel_size=1, - bias=False) + self.conv3 = build_conv_layer(self.conv_cfg, width, self.planes * self.expansion, kernel_size=1, bias=False) self.add_module(self.norm3_name, norm3) def forward(self, x): - def _inner_forward(x): identity = x @@ -285,16 +275,10 @@ class ResNeSt(ResNetV1d): 50: (Bottleneck, (3, 4, 6, 3)), 101: (Bottleneck, (3, 4, 23, 3)), 152: (Bottleneck, (3, 8, 36, 3)), - 200: (Bottleneck, (3, 24, 36, 3)) + 200: (Bottleneck, (3, 24, 36, 3)), } - def __init__(self, - groups=1, - base_width=4, - radix=2, - reduction_factor=4, - avg_down_stride=True, - **kwargs): + def __init__(self, groups=1, base_width=4, radix=2, reduction_factor=4, avg_down_stride=True, **kwargs): self.groups = groups self.base_width = base_width self.radix = radix @@ -311,4 +295,5 @@ def make_res_layer(self, **kwargs): radix=self.radix, reduction_factor=self.reduction_factor, avg_down_stride=self.avg_down_stride, - **kwargs) + **kwargs, + ) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnet.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnet.py index ffac52d2..03d26b45 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnet.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnet.py @@ -1,11 +1,16 @@ import torch.nn as nn import torch.utils.checkpoint as cp -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import (build_conv_layer, build_norm_layer, build_plugin_layer, - constant_init, kaiming_init) -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _BatchNorm - -from comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ( + build_conv_layer, + build_norm_layer, + build_plugin_layer, + constant_init, + kaiming_init, +) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _BatchNorm + +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger from ..builder import BACKBONES from ..utils import ResLayer @@ -15,37 +20,32 @@ class BasicBlock(nn.Module): expansion = 1 - def __init__(self, - inplanes, - planes, - stride=1, - dilation=1, - downsample=None, - style='pytorch', - with_cp=False, - conv_cfg=None, - norm_cfg=dict(type='BN'), - dcn=None, - plugins=None): + def __init__( + self, + inplanes, + planes, + stride=1, + dilation=1, + downsample=None, + style="pytorch", + with_cp=False, + conv_cfg=None, + norm_cfg=dict(type="BN"), + dcn=None, + plugins=None, + ): super(BasicBlock, self).__init__() - assert dcn is None, 'Not implemented yet.' - assert plugins is None, 'Not implemented yet.' + assert dcn is None, "Not implemented yet." + assert plugins is None, "Not implemented yet." self.norm1_name, norm1 = build_norm_layer(norm_cfg, planes, postfix=1) self.norm2_name, norm2 = build_norm_layer(norm_cfg, planes, postfix=2) self.conv1 = build_conv_layer( - conv_cfg, - inplanes, - planes, - 3, - stride=stride, - padding=dilation, - dilation=dilation, - bias=False) + conv_cfg, inplanes, planes, 3, stride=stride, padding=dilation, dilation=dilation, bias=False + ) self.add_module(self.norm1_name, norm1) - self.conv2 = build_conv_layer( - conv_cfg, planes, planes, 3, padding=1, bias=False) + self.conv2 = build_conv_layer(conv_cfg, planes, planes, 3, padding=1, bias=False) self.add_module(self.norm2_name, norm2) self.relu = nn.ReLU(inplace=True) @@ -103,25 +103,27 @@ class Bottleneck(nn.Module): expansion = 4 - def __init__(self, - inplanes, - planes, - stride=1, - dilation=1, - downsample=None, - style='pytorch', - with_cp=False, - conv_cfg=None, - norm_cfg=dict(type='BN'), - dcn=None, - plugins=None): + def __init__( + self, + inplanes, + planes, + stride=1, + dilation=1, + downsample=None, + style="pytorch", + with_cp=False, + conv_cfg=None, + norm_cfg=dict(type="BN"), + dcn=None, + plugins=None, + ): super(Bottleneck, self).__init__() - assert style in ['pytorch', 'caffe'] + assert style in ["pytorch", "caffe"] assert dcn is None or isinstance(dcn, dict) assert plugins is None or isinstance(plugins, list) if plugins is not None: - allowed_position = ['after_conv1', 'after_conv2', 'after_conv3'] - assert all(p['position'] in allowed_position for p in plugins) + allowed_position = ["after_conv1", "after_conv2", "after_conv3"] + assert all(p["position"] in allowed_position for p in plugins) self.inplanes = inplanes self.planes = planes @@ -138,20 +140,11 @@ def __init__(self, if self.with_plugins: # collect plugins for conv1/conv2/conv3 - self.after_conv1_plugins = [ - plugin['cfg'] for plugin in plugins - if plugin['position'] == 'after_conv1' - ] - self.after_conv2_plugins = [ - plugin['cfg'] for plugin in plugins - if plugin['position'] == 'after_conv2' - ] - self.after_conv3_plugins = [ - plugin['cfg'] for plugin in plugins - if plugin['position'] == 'after_conv3' - ] - - if self.style == 'pytorch': + self.after_conv1_plugins = [plugin["cfg"] for plugin in plugins if plugin["position"] == "after_conv1"] + self.after_conv2_plugins = [plugin["cfg"] for plugin in plugins if plugin["position"] == "after_conv2"] + self.after_conv3_plugins = [plugin["cfg"] for plugin in plugins if plugin["position"] == "after_conv3"] + + if self.style == "pytorch": self.conv1_stride = 1 self.conv2_stride = stride else: @@ -160,20 +153,13 @@ def __init__(self, self.norm1_name, norm1 = build_norm_layer(norm_cfg, planes, postfix=1) self.norm2_name, norm2 = build_norm_layer(norm_cfg, planes, postfix=2) - self.norm3_name, norm3 = build_norm_layer( - norm_cfg, planes * self.expansion, postfix=3) + self.norm3_name, norm3 = build_norm_layer(norm_cfg, planes * self.expansion, postfix=3) - self.conv1 = build_conv_layer( - conv_cfg, - inplanes, - planes, - kernel_size=1, - stride=self.conv1_stride, - bias=False) + self.conv1 = build_conv_layer(conv_cfg, inplanes, planes, kernel_size=1, stride=self.conv1_stride, bias=False) self.add_module(self.norm1_name, norm1) fallback_on_stride = False if self.with_dcn: - fallback_on_stride = dcn.pop('fallback_on_stride', False) + fallback_on_stride = dcn.pop("fallback_on_stride", False) if not self.with_dcn or fallback_on_stride: self.conv2 = build_conv_layer( conv_cfg, @@ -183,9 +169,10 @@ def __init__(self, stride=self.conv2_stride, padding=dilation, dilation=dilation, - bias=False) + bias=False, + ) else: - assert self.conv_cfg is None, 'conv_cfg must be None for DCN' + assert self.conv_cfg is None, "conv_cfg must be None for DCN" self.conv2 = build_conv_layer( dcn, planes, @@ -194,27 +181,20 @@ def __init__(self, stride=self.conv2_stride, padding=dilation, dilation=dilation, - bias=False) + bias=False, + ) self.add_module(self.norm2_name, norm2) - self.conv3 = build_conv_layer( - conv_cfg, - planes, - planes * self.expansion, - kernel_size=1, - bias=False) + self.conv3 = build_conv_layer(conv_cfg, planes, planes * self.expansion, kernel_size=1, bias=False) self.add_module(self.norm3_name, norm3) self.relu = nn.ReLU(inplace=True) self.downsample = downsample if self.with_plugins: - self.after_conv1_plugin_names = self.make_block_plugins( - planes, self.after_conv1_plugins) - self.after_conv2_plugin_names = self.make_block_plugins( - planes, self.after_conv2_plugins) - self.after_conv3_plugin_names = self.make_block_plugins( - planes * self.expansion, self.after_conv3_plugins) + self.after_conv1_plugin_names = self.make_block_plugins(planes, self.after_conv1_plugins) + self.after_conv2_plugin_names = self.make_block_plugins(planes, self.after_conv2_plugins) + self.after_conv3_plugin_names = self.make_block_plugins(planes * self.expansion, self.after_conv3_plugins) def make_block_plugins(self, in_channels, plugins): """make plugins for block. @@ -230,11 +210,8 @@ def make_block_plugins(self, in_channels, plugins): plugin_names = [] for plugin in plugins: plugin = plugin.copy() - name, layer = build_plugin_layer( - plugin, - in_channels=in_channels, - postfix=plugin.pop('postfix', '')) - assert not hasattr(self, name), f'duplicate plugin {name}' + name, layer = build_plugin_layer(plugin, in_channels=in_channels, postfix=plugin.pop("postfix", "")) + assert not hasattr(self, name), f"duplicate plugin {name}" self.add_module(name, layer) plugin_names.append(name) return plugin_names @@ -348,7 +325,7 @@ class ResNet(nn.Module): in resblocks to let them behave as identity. Example: - >>> from comfy_controlnet_preprocessors.uniformer.mmseg.models import ResNet + >>> from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.models import ResNet >>> import torch >>> self = ResNet(depth=18) >>> self.eval() @@ -367,35 +344,37 @@ class ResNet(nn.Module): 34: (BasicBlock, (3, 4, 6, 3)), 50: (Bottleneck, (3, 4, 6, 3)), 101: (Bottleneck, (3, 4, 23, 3)), - 152: (Bottleneck, (3, 8, 36, 3)) + 152: (Bottleneck, (3, 8, 36, 3)), } - def __init__(self, - depth, - in_channels=3, - stem_channels=64, - base_channels=64, - num_stages=4, - strides=(1, 2, 2, 2), - dilations=(1, 1, 1, 1), - out_indices=(0, 1, 2, 3), - style='pytorch', - deep_stem=False, - avg_down=False, - frozen_stages=-1, - conv_cfg=None, - norm_cfg=dict(type='BN', requires_grad=True), - norm_eval=False, - dcn=None, - stage_with_dcn=(False, False, False, False), - plugins=None, - multi_grid=None, - contract_dilation=False, - with_cp=False, - zero_init_residual=True): + def __init__( + self, + depth, + in_channels=3, + stem_channels=64, + base_channels=64, + num_stages=4, + strides=(1, 2, 2, 2), + dilations=(1, 1, 1, 1), + out_indices=(0, 1, 2, 3), + style="pytorch", + deep_stem=False, + avg_down=False, + frozen_stages=-1, + conv_cfg=None, + norm_cfg=dict(type="BN", requires_grad=True), + norm_eval=False, + dcn=None, + stage_with_dcn=(False, False, False, False), + plugins=None, + multi_grid=None, + contract_dilation=False, + with_cp=False, + zero_init_residual=True, + ): super(ResNet, self).__init__() if depth not in self.arch_settings: - raise KeyError(f'invalid depth {depth} for resnet') + raise KeyError(f"invalid depth {depth} for resnet") self.depth = depth self.stem_channels = stem_channels self.base_channels = base_channels @@ -438,8 +417,7 @@ def __init__(self, else: stage_plugins = None # multi grid is applied to last layer only - stage_multi_grid = multi_grid if i == len( - self.stage_blocks) - 1 else None + stage_multi_grid = multi_grid if i == len(self.stage_blocks) - 1 else None planes = base_channels * 2**i res_layer = self.make_res_layer( block=self.block, @@ -456,16 +434,16 @@ def __init__(self, dcn=dcn, plugins=stage_plugins, multi_grid=stage_multi_grid, - contract_dilation=contract_dilation) + contract_dilation=contract_dilation, + ) self.inplanes = planes * self.block.expansion - layer_name = f'layer{i+1}' + layer_name = f"layer{i+1}" self.add_module(layer_name, res_layer) self.res_layers.append(layer_name) self._freeze_stages() - self.feat_dim = self.block.expansion * base_channels * 2**( - len(self.stage_blocks) - 1) + self.feat_dim = self.block.expansion * base_channels * 2 ** (len(self.stage_blocks) - 1) def make_stage_plugins(self, plugins, stage_idx): """make plugins for ResNet 'stage_idx'th stage . @@ -512,7 +490,7 @@ def make_stage_plugins(self, plugins, stage_idx): stage_plugins = [] for plugin in plugins: plugin = plugin.copy() - stages = plugin.pop('stages', None) + stages = plugin.pop("stages", None) assert stages is None or len(stages) == self.num_stages # whether to insert plugin into current stage if stages is None or stages[stage_idx]: @@ -534,13 +512,8 @@ def _make_stem_layer(self, in_channels, stem_channels): if self.deep_stem: self.stem = nn.Sequential( build_conv_layer( - self.conv_cfg, - in_channels, - stem_channels // 2, - kernel_size=3, - stride=2, - padding=1, - bias=False), + self.conv_cfg, in_channels, stem_channels // 2, kernel_size=3, stride=2, padding=1, bias=False + ), build_norm_layer(self.norm_cfg, stem_channels // 2)[1], nn.ReLU(inplace=True), build_conv_layer( @@ -550,30 +523,21 @@ def _make_stem_layer(self, in_channels, stem_channels): kernel_size=3, stride=1, padding=1, - bias=False), + bias=False, + ), build_norm_layer(self.norm_cfg, stem_channels // 2)[1], nn.ReLU(inplace=True), build_conv_layer( - self.conv_cfg, - stem_channels // 2, - stem_channels, - kernel_size=3, - stride=1, - padding=1, - bias=False), + self.conv_cfg, stem_channels // 2, stem_channels, kernel_size=3, stride=1, padding=1, bias=False + ), build_norm_layer(self.norm_cfg, stem_channels)[1], - nn.ReLU(inplace=True)) + nn.ReLU(inplace=True), + ) else: self.conv1 = build_conv_layer( - self.conv_cfg, - in_channels, - stem_channels, - kernel_size=7, - stride=2, - padding=3, - bias=False) - self.norm1_name, norm1 = build_norm_layer( - self.norm_cfg, stem_channels, postfix=1) + self.conv_cfg, in_channels, stem_channels, kernel_size=7, stride=2, padding=3, bias=False + ) + self.norm1_name, norm1 = build_norm_layer(self.norm_cfg, stem_channels, postfix=1) self.add_module(self.norm1_name, norm1) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) @@ -592,7 +556,7 @@ def _freeze_stages(self): param.requires_grad = False for i in range(1, self.frozen_stages + 1): - m = getattr(self, f'layer{i}') + m = getattr(self, f"layer{i}") m.eval() for param in m.parameters(): param.requires_grad = False @@ -616,8 +580,7 @@ def init_weights(self, pretrained=None): if self.dcn is not None: for m in self.modules(): - if isinstance(m, Bottleneck) and hasattr( - m, 'conv2_offset'): + if isinstance(m, Bottleneck) and hasattr(m, "conv2_offset"): constant_init(m.conv2_offset, 0) if self.zero_init_residual: @@ -627,7 +590,7 @@ def init_weights(self, pretrained=None): elif isinstance(m, BasicBlock): constant_init(m.norm2, 0) else: - raise TypeError('pretrained must be a str or None') + raise TypeError("pretrained must be a str or None") def forward(self, x): """Forward function.""" @@ -670,8 +633,7 @@ class ResNetV1c(ResNet): """ def __init__(self, **kwargs): - super(ResNetV1c, self).__init__( - deep_stem=True, avg_down=False, **kwargs) + super(ResNetV1c, self).__init__(deep_stem=True, avg_down=False, **kwargs) @BACKBONES.register_module() @@ -684,5 +646,4 @@ class ResNetV1d(ResNet): """ def __init__(self, **kwargs): - super(ResNetV1d, self).__init__( - deep_stem=True, avg_down=True, **kwargs) + super(ResNetV1d, self).__init__(deep_stem=True, avg_down=True, **kwargs) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnext.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnext.py index f0a2d296..f9071c0b 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnext.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/resnext.py @@ -1,6 +1,6 @@ import math -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import build_conv_layer, build_norm_layer +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import build_conv_layer, build_norm_layer from ..builder import BACKBONES from ..utils import ResLayer @@ -15,40 +15,26 @@ class Bottleneck(_Bottleneck): "caffe", the stride-two layer is the first 1x1 conv layer. """ - def __init__(self, - inplanes, - planes, - groups=1, - base_width=4, - base_channels=64, - **kwargs): + def __init__(self, inplanes, planes, groups=1, base_width=4, base_channels=64, **kwargs): super(Bottleneck, self).__init__(inplanes, planes, **kwargs) if groups == 1: width = self.planes else: - width = math.floor(self.planes * - (base_width / base_channels)) * groups + width = math.floor(self.planes * (base_width / base_channels)) * groups - self.norm1_name, norm1 = build_norm_layer( - self.norm_cfg, width, postfix=1) - self.norm2_name, norm2 = build_norm_layer( - self.norm_cfg, width, postfix=2) - self.norm3_name, norm3 = build_norm_layer( - self.norm_cfg, self.planes * self.expansion, postfix=3) + self.norm1_name, norm1 = build_norm_layer(self.norm_cfg, width, postfix=1) + self.norm2_name, norm2 = build_norm_layer(self.norm_cfg, width, postfix=2) + self.norm3_name, norm3 = build_norm_layer(self.norm_cfg, self.planes * self.expansion, postfix=3) self.conv1 = build_conv_layer( - self.conv_cfg, - self.inplanes, - width, - kernel_size=1, - stride=self.conv1_stride, - bias=False) + self.conv_cfg, self.inplanes, width, kernel_size=1, stride=self.conv1_stride, bias=False + ) self.add_module(self.norm1_name, norm1) fallback_on_stride = False self.with_modulated_dcn = False if self.with_dcn: - fallback_on_stride = self.dcn.pop('fallback_on_stride', False) + fallback_on_stride = self.dcn.pop("fallback_on_stride", False) if not self.with_dcn or fallback_on_stride: self.conv2 = build_conv_layer( self.conv_cfg, @@ -59,9 +45,10 @@ def __init__(self, padding=self.dilation, dilation=self.dilation, groups=groups, - bias=False) + bias=False, + ) else: - assert self.conv_cfg is None, 'conv_cfg must be None for DCN' + assert self.conv_cfg is None, "conv_cfg must be None for DCN" self.conv2 = build_conv_layer( self.dcn, width, @@ -71,15 +58,11 @@ def __init__(self, padding=self.dilation, dilation=self.dilation, groups=groups, - bias=False) + bias=False, + ) self.add_module(self.norm2_name, norm2) - self.conv3 = build_conv_layer( - self.conv_cfg, - width, - self.planes * self.expansion, - kernel_size=1, - bias=False) + self.conv3 = build_conv_layer(self.conv_cfg, width, self.planes * self.expansion, kernel_size=1, bias=False) self.add_module(self.norm3_name, norm3) @@ -111,7 +94,7 @@ class ResNeXt(ResNet): in resblocks to let them behave as identity. Example: - >>> from comfy_controlnet_preprocessors.uniformer.mmseg.models import ResNeXt + >>> from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.models import ResNeXt >>> import torch >>> self = ResNeXt(depth=50) >>> self.eval() @@ -128,7 +111,7 @@ class ResNeXt(ResNet): arch_settings = { 50: (Bottleneck, (3, 4, 6, 3)), 101: (Bottleneck, (3, 4, 23, 3)), - 152: (Bottleneck, (3, 8, 36, 3)) + 152: (Bottleneck, (3, 8, 36, 3)), } def __init__(self, groups=1, base_width=4, **kwargs): @@ -138,8 +121,4 @@ def __init__(self, groups=1, base_width=4, **kwargs): def make_res_layer(self, **kwargs): """Pack all blocks in a stage into a ``ResLayer``""" - return ResLayer( - groups=self.groups, - base_width=self.base_width, - base_channels=self.base_channels, - **kwargs) + return ResLayer(groups=self.groups, base_width=self.base_width, base_channels=self.base_channels, **kwargs) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/unet.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/unet.py index 68be690f..57b648a3 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/unet.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/unet.py @@ -1,11 +1,17 @@ import torch.nn as nn import torch.utils.checkpoint as cp -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import (UPSAMPLE_LAYERS, ConvModule, build_activation_layer, - build_norm_layer, constant_init, kaiming_init) -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _BatchNorm - -from comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ( + UPSAMPLE_LAYERS, + ConvModule, + build_activation_layer, + build_norm_layer, + constant_init, + kaiming_init, +) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import load_checkpoint +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _BatchNorm + +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger from ..builder import BACKBONES from ..utils import UpConvBlock @@ -40,21 +46,23 @@ class BasicConvBlock(nn.Module): plugins (dict): plugins for convolutional layers. Default: None. """ - def __init__(self, - in_channels, - out_channels, - num_convs=2, - stride=1, - dilation=1, - with_cp=False, - conv_cfg=None, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU'), - dcn=None, - plugins=None): + def __init__( + self, + in_channels, + out_channels, + num_convs=2, + stride=1, + dilation=1, + with_cp=False, + conv_cfg=None, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU"), + dcn=None, + plugins=None, + ): super(BasicConvBlock, self).__init__() - assert dcn is None, 'Not implemented yet.' - assert plugins is None, 'Not implemented yet.' + assert dcn is None, "Not implemented yet." + assert plugins is None, "Not implemented yet." self.with_cp = with_cp convs = [] @@ -69,7 +77,9 @@ def __init__(self, padding=1 if i == 0 else dilation, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg)) + act_cfg=act_cfg, + ) + ) self.convs = nn.Sequential(*convs) @@ -102,33 +112,30 @@ class DeconvModule(nn.Module): kernel_size (int): Kernel size of the convolutional layer. Default: 4. """ - def __init__(self, - in_channels, - out_channels, - with_cp=False, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU'), - *, - kernel_size=4, - scale_factor=2): + def __init__( + self, + in_channels, + out_channels, + with_cp=False, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU"), + *, + kernel_size=4, + scale_factor=2, + ): super(DeconvModule, self).__init__() - assert (kernel_size - scale_factor >= 0) and\ - (kernel_size - scale_factor) % 2 == 0,\ - f'kernel_size should be greater than or equal to scale_factor '\ - f'and (kernel_size - scale_factor) should be even numbers, '\ - f'while the kernel size is {kernel_size} and scale_factor is '\ - f'{scale_factor}.' + assert (kernel_size - scale_factor >= 0) and (kernel_size - scale_factor) % 2 == 0, ( + f"kernel_size should be greater than or equal to scale_factor " + f"and (kernel_size - scale_factor) should be even numbers, " + f"while the kernel size is {kernel_size} and scale_factor is " + f"{scale_factor}." + ) stride = scale_factor padding = (kernel_size - scale_factor) // 2 self.with_cp = with_cp - deconv = nn.ConvTranspose2d( - in_channels, - out_channels, - kernel_size=kernel_size, - stride=stride, - padding=padding) + deconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding) norm_name, norm = build_norm_layer(norm_cfg, out_channels) activate = build_activation_layer(act_cfg) @@ -176,20 +183,21 @@ class InterpConv(nn.Module): scale_factor=2, mode='bilinear', align_corners=False). """ - def __init__(self, - in_channels, - out_channels, - with_cp=False, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU'), - *, - conv_cfg=None, - conv_first=False, - kernel_size=1, - stride=1, - padding=0, - upsample_cfg=dict( - scale_factor=2, mode='bilinear', align_corners=False)): + def __init__( + self, + in_channels, + out_channels, + with_cp=False, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU"), + *, + conv_cfg=None, + conv_first=False, + kernel_size=1, + stride=1, + padding=0, + upsample_cfg=dict(scale_factor=2, mode="bilinear", align_corners=False), + ): super(InterpConv, self).__init__() self.with_cp = with_cp @@ -201,7 +209,8 @@ def __init__(self, padding=padding, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) upsample = nn.Upsample(**upsample_cfg) if conv_first: self.interp_upsample = nn.Sequential(conv, upsample) @@ -274,57 +283,65 @@ class UNet(nn.Module): """ - def __init__(self, - in_channels=3, - base_channels=64, - num_stages=5, - strides=(1, 1, 1, 1, 1), - enc_num_convs=(2, 2, 2, 2, 2), - dec_num_convs=(2, 2, 2, 2), - downsamples=(True, True, True, True), - enc_dilations=(1, 1, 1, 1, 1), - dec_dilations=(1, 1, 1, 1), - with_cp=False, - conv_cfg=None, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU'), - upsample_cfg=dict(type='InterpConv'), - norm_eval=False, - dcn=None, - plugins=None): + def __init__( + self, + in_channels=3, + base_channels=64, + num_stages=5, + strides=(1, 1, 1, 1, 1), + enc_num_convs=(2, 2, 2, 2, 2), + dec_num_convs=(2, 2, 2, 2), + downsamples=(True, True, True, True), + enc_dilations=(1, 1, 1, 1, 1), + dec_dilations=(1, 1, 1, 1), + with_cp=False, + conv_cfg=None, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU"), + upsample_cfg=dict(type="InterpConv"), + norm_eval=False, + dcn=None, + plugins=None, + ): super(UNet, self).__init__() - assert dcn is None, 'Not implemented yet.' - assert plugins is None, 'Not implemented yet.' - assert len(strides) == num_stages, \ - 'The length of strides should be equal to num_stages, '\ - f'while the strides is {strides}, the length of '\ - f'strides is {len(strides)}, and the num_stages is '\ - f'{num_stages}.' - assert len(enc_num_convs) == num_stages, \ - 'The length of enc_num_convs should be equal to num_stages, '\ - f'while the enc_num_convs is {enc_num_convs}, the length of '\ - f'enc_num_convs is {len(enc_num_convs)}, and the num_stages is '\ - f'{num_stages}.' - assert len(dec_num_convs) == (num_stages-1), \ - 'The length of dec_num_convs should be equal to (num_stages-1), '\ - f'while the dec_num_convs is {dec_num_convs}, the length of '\ - f'dec_num_convs is {len(dec_num_convs)}, and the num_stages is '\ - f'{num_stages}.' - assert len(downsamples) == (num_stages-1), \ - 'The length of downsamples should be equal to (num_stages-1), '\ - f'while the downsamples is {downsamples}, the length of '\ - f'downsamples is {len(downsamples)}, and the num_stages is '\ - f'{num_stages}.' - assert len(enc_dilations) == num_stages, \ - 'The length of enc_dilations should be equal to num_stages, '\ - f'while the enc_dilations is {enc_dilations}, the length of '\ - f'enc_dilations is {len(enc_dilations)}, and the num_stages is '\ - f'{num_stages}.' - assert len(dec_dilations) == (num_stages-1), \ - 'The length of dec_dilations should be equal to (num_stages-1), '\ - f'while the dec_dilations is {dec_dilations}, the length of '\ - f'dec_dilations is {len(dec_dilations)}, and the num_stages is '\ - f'{num_stages}.' + assert dcn is None, "Not implemented yet." + assert plugins is None, "Not implemented yet." + assert len(strides) == num_stages, ( + "The length of strides should be equal to num_stages, " + f"while the strides is {strides}, the length of " + f"strides is {len(strides)}, and the num_stages is " + f"{num_stages}." + ) + assert len(enc_num_convs) == num_stages, ( + "The length of enc_num_convs should be equal to num_stages, " + f"while the enc_num_convs is {enc_num_convs}, the length of " + f"enc_num_convs is {len(enc_num_convs)}, and the num_stages is " + f"{num_stages}." + ) + assert len(dec_num_convs) == (num_stages - 1), ( + "The length of dec_num_convs should be equal to (num_stages-1), " + f"while the dec_num_convs is {dec_num_convs}, the length of " + f"dec_num_convs is {len(dec_num_convs)}, and the num_stages is " + f"{num_stages}." + ) + assert len(downsamples) == (num_stages - 1), ( + "The length of downsamples should be equal to (num_stages-1), " + f"while the downsamples is {downsamples}, the length of " + f"downsamples is {len(downsamples)}, and the num_stages is " + f"{num_stages}." + ) + assert len(enc_dilations) == num_stages, ( + "The length of enc_dilations should be equal to num_stages, " + f"while the enc_dilations is {enc_dilations}, the length of " + f"enc_dilations is {len(enc_dilations)}, and the num_stages is " + f"{num_stages}." + ) + assert len(dec_dilations) == (num_stages - 1), ( + "The length of dec_dilations should be equal to (num_stages-1), " + f"while the dec_dilations is {dec_dilations}, the length of " + f"dec_dilations is {len(dec_dilations)}, and the num_stages is " + f"{num_stages}." + ) self.num_stages = num_stages self.strides = strides self.downsamples = downsamples @@ -339,13 +356,13 @@ def __init__(self, if i != 0: if strides[i] == 1 and downsamples[i - 1]: enc_conv_block.append(nn.MaxPool2d(kernel_size=2)) - upsample = (strides[i] != 1 or downsamples[i - 1]) + upsample = strides[i] != 1 or downsamples[i - 1] self.decoder.append( UpConvBlock( conv_block=BasicConvBlock, in_channels=base_channels * 2**i, - skip_channels=base_channels * 2**(i - 1), - out_channels=base_channels * 2**(i - 1), + skip_channels=base_channels * 2 ** (i - 1), + out_channels=base_channels * 2 ** (i - 1), num_convs=dec_num_convs[i - 1], stride=1, dilation=dec_dilations[i - 1], @@ -355,7 +372,9 @@ def __init__(self, act_cfg=act_cfg, upsample_cfg=upsample_cfg if upsample else None, dcn=None, - plugins=None)) + plugins=None, + ) + ) enc_conv_block.append( BasicConvBlock( @@ -369,7 +388,9 @@ def __init__(self, norm_cfg=norm_cfg, act_cfg=act_cfg, dcn=None, - plugins=None)) + plugins=None, + ) + ) self.encoder.append((nn.Sequential(*enc_conv_block))) in_channels = base_channels * 2**i @@ -402,12 +423,12 @@ def _check_input_divisible(self, x): for i in range(1, self.num_stages): if self.strides[i] == 2 or self.downsamples[i - 1]: whole_downsample_rate *= 2 - assert (h % whole_downsample_rate == 0) \ - and (w % whole_downsample_rate == 0),\ - f'The input image size {(h, w)} should be divisible by the whole '\ - f'downsample rate {whole_downsample_rate}, when num_stages is '\ - f'{self.num_stages}, strides is {self.strides}, and downsamples '\ - f'is {self.downsamples}.' + assert (h % whole_downsample_rate == 0) and (w % whole_downsample_rate == 0), ( + f"The input image size {(h, w)} should be divisible by the whole " + f"downsample rate {whole_downsample_rate}, when num_stages is " + f"{self.num_stages}, strides is {self.strides}, and downsamples " + f"is {self.downsamples}." + ) def init_weights(self, pretrained=None): """Initialize the weights in backbone. @@ -426,4 +447,4 @@ def init_weights(self, pretrained=None): elif isinstance(m, (_BatchNorm, nn.GroupNorm)): constant_init(m, 1) else: - raise TypeError('pretrained must be a str or None') + raise TypeError("pretrained must be a str or None") diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/uniformer.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/uniformer.py index 3b808710..2f09766f 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/uniformer.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/uniformer.py @@ -16,13 +16,13 @@ import numpy as np from timm.models.layers import DropPath, to_2tuple, trunc_normal_ -from comfy_controlnet_preprocessors.uniformer.mmcv_custom import load_checkpoint -from comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv_custom import load_checkpoint +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger from ..builder import BACKBONES class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -41,7 +41,7 @@ def forward(self, x): class CMlp(nn.Module): - def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -60,8 +60,19 @@ def forward(self, x): class CBlock(nn.Module): - def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): + def __init__( + self, + dim, + num_heads, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + ): super().__init__() self.pos_embed = nn.Conv2d(dim, dim, 3, padding=1, groups=dim) self.norm1 = nn.BatchNorm2d(dim) @@ -69,7 +80,7 @@ def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, self.conv2 = nn.Conv2d(dim, dim, 1) self.attn = nn.Conv2d(dim, dim, 5, padding=2, groups=dim) # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = nn.BatchNorm2d(dim) mlp_hidden_dim = int(dim * mlp_ratio) self.mlp = CMlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) @@ -82,12 +93,12 @@ def forward(self, x): class Attention(nn.Module): - def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0): super().__init__() self.num_heads = num_heads head_dim = dim // num_heads # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights - self.scale = qk_scale or head_dim ** -0.5 + self.scale = qk_scale or head_dim**-0.5 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) @@ -97,7 +108,7 @@ def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0. def forward(self, x): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) - q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) attn = (q @ k.transpose(-2, -1)) * self.scale attn = attn.softmax(dim=-1) @@ -110,17 +121,27 @@ def forward(self, x): class SABlock(nn.Module): - def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): + def __init__( + self, + dim, + num_heads, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + ): super().__init__() self.pos_embed = nn.Conv2d(dim, dim, 3, padding=1, groups=dim) self.norm1 = norm_layer(dim) self.attn = Attention( - dim, - num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, - attn_drop=attn_drop, proj_drop=drop) + dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop + ) # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) @@ -166,18 +187,29 @@ def window_reverse(windows, window_size, H, W): class SABlock_Windows(nn.Module): - def __init__(self, dim, num_heads, window_size=14, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): + def __init__( + self, + dim, + num_heads, + window_size=14, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + ): super().__init__() - self.window_size=window_size + self.window_size = window_size self.pos_embed = nn.Conv2d(dim, dim, 3, padding=1, groups=dim) self.norm1 = norm_layer(dim) self.attn = Attention( - dim, - num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, - attn_drop=attn_drop, proj_drop=drop) + dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop + ) # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) @@ -216,8 +248,8 @@ def forward(self, x): class PatchEmbed(nn.Module): - """ Image to Patch Embedding - """ + """Image to Patch Embedding""" + def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): super().__init__() img_size = to_2tuple(img_size) @@ -241,15 +273,34 @@ def forward(self, x): @BACKBONES.register_module() class UniFormer(nn.Module): - """ Vision Transformer + """Vision Transformer A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale` - https://arxiv.org/abs/2010.11929 """ - def __init__(self, layers=[3, 4, 8, 3], img_size=224, in_chans=3, num_classes=80, embed_dim=[64, 128, 320, 512], - head_dim=64, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None, - drop_rate=0., attn_drop_rate=0., drop_path_rate=0., norm_layer=partial(nn.LayerNorm, eps=1e-6), - pretrained_path=None, use_checkpoint=False, checkpoint_num=[0, 0, 0, 0], - windows=False, hybrid=False, window_size=14): + + def __init__( + self, + layers=[3, 4, 8, 3], + img_size=224, + in_chans=3, + num_classes=80, + embed_dim=[64, 128, 320, 512], + head_dim=64, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + representation_size=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.0, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + pretrained_path=None, + use_checkpoint=False, + checkpoint_num=[0, 0, 0, 0], + windows=False, + hybrid=False, + window_size=14, + ): """ Args: layer (list): number of block in each layer @@ -278,77 +329,155 @@ def __init__(self, layers=[3, 4, 8, 3], img_size=224, in_chans=3, num_classes=80 self.use_checkpoint = use_checkpoint self.checkpoint_num = checkpoint_num self.windows = windows - print(f'Use Checkpoint: {self.use_checkpoint}') - print(f'Checkpoint Number: {self.checkpoint_num}') + print(f"Use Checkpoint: {self.use_checkpoint}") + print(f"Checkpoint Number: {self.checkpoint_num}") self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) - self.patch_embed1 = PatchEmbed( - img_size=img_size, patch_size=4, in_chans=in_chans, embed_dim=embed_dim[0]) + self.patch_embed1 = PatchEmbed(img_size=img_size, patch_size=4, in_chans=in_chans, embed_dim=embed_dim[0]) self.patch_embed2 = PatchEmbed( - img_size=img_size // 4, patch_size=2, in_chans=embed_dim[0], embed_dim=embed_dim[1]) + img_size=img_size // 4, patch_size=2, in_chans=embed_dim[0], embed_dim=embed_dim[1] + ) self.patch_embed3 = PatchEmbed( - img_size=img_size // 8, patch_size=2, in_chans=embed_dim[1], embed_dim=embed_dim[2]) + img_size=img_size // 8, patch_size=2, in_chans=embed_dim[1], embed_dim=embed_dim[2] + ) self.patch_embed4 = PatchEmbed( - img_size=img_size // 16, patch_size=2, in_chans=embed_dim[2], embed_dim=embed_dim[3]) + img_size=img_size // 16, patch_size=2, in_chans=embed_dim[2], embed_dim=embed_dim[3] + ) self.pos_drop = nn.Dropout(p=drop_rate) dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(layers))] # stochastic depth decay rule num_heads = [dim // head_dim for dim in embed_dim] - self.blocks1 = nn.ModuleList([ - CBlock( - dim=embed_dim[0], num_heads=num_heads[0], mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer) - for i in range(layers[0])]) - self.norm1=norm_layer(embed_dim[0]) - self.blocks2 = nn.ModuleList([ - CBlock( - dim=embed_dim[1], num_heads=num_heads[1], mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i+layers[0]], norm_layer=norm_layer) - for i in range(layers[1])]) + self.blocks1 = nn.ModuleList( + [ + CBlock( + dim=embed_dim[0], + num_heads=num_heads[0], + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + ) + for i in range(layers[0]) + ] + ) + self.norm1 = norm_layer(embed_dim[0]) + self.blocks2 = nn.ModuleList( + [ + CBlock( + dim=embed_dim[1], + num_heads=num_heads[1], + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i + layers[0]], + norm_layer=norm_layer, + ) + for i in range(layers[1]) + ] + ) self.norm2 = norm_layer(embed_dim[1]) if self.windows: - print('Use local window for all blocks in stage3') - self.blocks3 = nn.ModuleList([ - SABlock_Windows( - dim=embed_dim[2], num_heads=num_heads[2], window_size=window_size, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i+layers[0]+layers[1]], norm_layer=norm_layer) - for i in range(layers[2])]) + print("Use local window for all blocks in stage3") + self.blocks3 = nn.ModuleList( + [ + SABlock_Windows( + dim=embed_dim[2], + num_heads=num_heads[2], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i + layers[0] + layers[1]], + norm_layer=norm_layer, + ) + for i in range(layers[2]) + ] + ) elif hybrid: - print('Use hybrid window for blocks in stage3') + print("Use hybrid window for blocks in stage3") block3 = [] for i in range(layers[2]): if (i + 1) % 4 == 0: - block3.append(SABlock( - dim=embed_dim[2], num_heads=num_heads[2], mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i+layers[0]+layers[1]], norm_layer=norm_layer)) + block3.append( + SABlock( + dim=embed_dim[2], + num_heads=num_heads[2], + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i + layers[0] + layers[1]], + norm_layer=norm_layer, + ) + ) else: - block3.append(SABlock_Windows( - dim=embed_dim[2], num_heads=num_heads[2], window_size=window_size, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i+layers[0]+layers[1]], norm_layer=norm_layer)) + block3.append( + SABlock_Windows( + dim=embed_dim[2], + num_heads=num_heads[2], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i + layers[0] + layers[1]], + norm_layer=norm_layer, + ) + ) self.blocks3 = nn.ModuleList(block3) else: - print('Use global window for all blocks in stage3') - self.blocks3 = nn.ModuleList([ - SABlock( - dim=embed_dim[2], num_heads=num_heads[2], mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i+layers[0]+layers[1]], norm_layer=norm_layer) - for i in range(layers[2])]) + print("Use global window for all blocks in stage3") + self.blocks3 = nn.ModuleList( + [ + SABlock( + dim=embed_dim[2], + num_heads=num_heads[2], + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i + layers[0] + layers[1]], + norm_layer=norm_layer, + ) + for i in range(layers[2]) + ] + ) self.norm3 = norm_layer(embed_dim[2]) - self.blocks4 = nn.ModuleList([ - SABlock( - dim=embed_dim[3], num_heads=num_heads[3], mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i+layers[0]+layers[1]+layers[2]], norm_layer=norm_layer) - for i in range(layers[3])]) + self.blocks4 = nn.ModuleList( + [ + SABlock( + dim=embed_dim[3], + num_heads=num_heads[3], + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i + layers[0] + layers[1] + layers[2]], + norm_layer=norm_layer, + ) + for i in range(layers[3]) + ] + ) self.norm4 = norm_layer(embed_dim[3]) # Representation layer if representation_size: self.num_features = representation_size - self.pre_logits = nn.Sequential(OrderedDict([ - ('fc', nn.Linear(embed_dim, representation_size)), - ('act', nn.Tanh()) - ])) + self.pre_logits = nn.Sequential( + OrderedDict([("fc", nn.Linear(embed_dim, representation_size)), ("act", nn.Tanh())]) + ) else: self.pre_logits = nn.Identity() @@ -358,11 +487,12 @@ def __init__(self, layers=[3, 4, 8, 3], img_size=224, in_chans=3, num_classes=80 def init_weights(self, pretrained): if isinstance(pretrained, str): logger = get_root_logger() - load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger) - print(f'Load pretrained model from {pretrained}') + load_checkpoint(self, pretrained, map_location="cpu", strict=False, logger=logger) + print(f"Load pretrained model from {pretrained}") + def _init_weights(self, m): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -371,12 +501,12 @@ def _init_weights(self, m): @torch.jit.ignore def no_weight_decay(self): - return {'pos_embed', 'cls_token'} + return {"pos_embed", "cls_token"} def get_classifier(self): return self.head - def reset_classifier(self, num_classes, global_pool=''): + def reset_classifier(self, num_classes, global_pool=""): self.num_classes = num_classes self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/vit.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/vit.py index 820fccf3..9c033551 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/vit.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/backbones/vit.py @@ -7,12 +7,19 @@ import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint as cp -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import (Conv2d, Linear, build_activation_layer, build_norm_layer, - constant_init, kaiming_init, normal_init) -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import _load_checkpoint -from comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _BatchNorm - -from comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ( + Conv2d, + Linear, + build_activation_layer, + build_norm_layer, + constant_init, + kaiming_init, + normal_init, +) +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import _load_checkpoint +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils.parrots_wrapper import _BatchNorm + +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.utils import get_root_logger from ..builder import BACKBONES from ..utils import DropPath, trunc_normal_ @@ -33,12 +40,7 @@ class Mlp(nn.Module): to be between 0 and 1. Default: 0. """ - def __init__(self, - in_features, - hidden_features=None, - out_features=None, - act_cfg=dict(type='GELU'), - drop=0.): + def __init__(self, in_features, hidden_features=None, out_features=None, act_cfg=dict(type="GELU"), drop=0.0): super(Mlp, self).__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -69,13 +71,7 @@ class Attention(nn.Module): proj_drop (float): Drop rate for output weights. Default: 0. """ - def __init__(self, - dim, - num_heads=8, - qkv_bias=False, - qk_scale=None, - attn_drop=0., - proj_drop=0.): + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0): super(Attention, self).__init__() self.num_heads = num_heads head_dim = dim // num_heads @@ -88,8 +84,7 @@ def __init__(self, def forward(self, x): b, n, c = x.shape - qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, - c // self.num_heads).permute(2, 0, 3, 1, 4) + qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, c // self.num_heads).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] attn = (q @ k.transpose(-2, -1)) * self.scale @@ -125,36 +120,31 @@ class Block(nn.Module): memory while slowing down the training speed. Default: False. """ - def __init__(self, - dim, - num_heads, - mlp_ratio=4, - qkv_bias=False, - qk_scale=None, - drop=0., - attn_drop=0., - proj_drop=0., - drop_path=0., - act_cfg=dict(type='GELU'), - norm_cfg=dict(type='LN', eps=1e-6), - with_cp=False): + def __init__( + self, + dim, + num_heads, + mlp_ratio=4, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + proj_drop=0.0, + drop_path=0.0, + act_cfg=dict(type="GELU"), + norm_cfg=dict(type="LN", eps=1e-6), + with_cp=False, + ): super(Block, self).__init__() self.with_cp = with_cp _, self.norm1 = build_norm_layer(norm_cfg, dim) - self.attn = Attention(dim, num_heads, qkv_bias, qk_scale, attn_drop, - proj_drop) - self.drop_path = DropPath( - drop_path) if drop_path > 0. else nn.Identity() + self.attn = Attention(dim, num_heads, qkv_bias, qk_scale, attn_drop, proj_drop) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() _, self.norm2 = build_norm_layer(norm_cfg, dim) mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp( - in_features=dim, - hidden_features=mlp_hidden_dim, - act_cfg=act_cfg, - drop=drop) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_cfg=act_cfg, drop=drop) def forward(self, x): - def _inner_forward(x): out = x + self.drop_path(self.attn(self.norm1(x))) out = out + self.drop_path(self.mlp(self.norm2(out))) @@ -180,23 +170,18 @@ class PatchEmbed(nn.Module): embed_dim (int): The embedding dimension. Default: 768. """ - def __init__(self, - img_size=224, - patch_size=16, - in_channels=3, - embed_dim=768): + def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768): super(PatchEmbed, self).__init__() if isinstance(img_size, int): self.img_size = (img_size, img_size) elif isinstance(img_size, tuple): self.img_size = img_size else: - raise TypeError('img_size must be type of int or tuple') + raise TypeError("img_size must be type of int or tuple") h, w = self.img_size self.patch_size = (patch_size, patch_size) self.num_patches = (h // patch_size) * (w // patch_size) - self.proj = Conv2d( - in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) + self.proj = Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): return self.proj(x).flatten(2).transpose(1, 2) @@ -243,41 +228,40 @@ class VisionTransformer(nn.Module): Default: False. """ - def __init__(self, - img_size=(224, 224), - patch_size=16, - in_channels=3, - embed_dim=768, - depth=12, - num_heads=12, - mlp_ratio=4, - out_indices=11, - qkv_bias=True, - qk_scale=None, - drop_rate=0., - attn_drop_rate=0., - drop_path_rate=0., - norm_cfg=dict(type='LN', eps=1e-6, requires_grad=True), - act_cfg=dict(type='GELU'), - norm_eval=False, - final_norm=False, - with_cls_token=True, - interpolate_mode='bicubic', - with_cp=False): + def __init__( + self, + img_size=(224, 224), + patch_size=16, + in_channels=3, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4, + out_indices=11, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.0, + norm_cfg=dict(type="LN", eps=1e-6, requires_grad=True), + act_cfg=dict(type="GELU"), + norm_eval=False, + final_norm=False, + with_cls_token=True, + interpolate_mode="bicubic", + with_cp=False, + ): super(VisionTransformer, self).__init__() self.img_size = img_size self.patch_size = patch_size self.features = self.embed_dim = embed_dim self.patch_embed = PatchEmbed( - img_size=img_size, - patch_size=patch_size, - in_channels=in_channels, - embed_dim=embed_dim) + img_size=img_size, patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim + ) self.with_cls_token = with_cls_token self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim)) - self.pos_embed = nn.Parameter( - torch.zeros(1, self.patch_embed.num_patches + 1, embed_dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_embed.num_patches + 1, embed_dim)) self.pos_drop = nn.Dropout(p=drop_rate) if isinstance(out_indices, int): @@ -285,23 +269,26 @@ def __init__(self, elif isinstance(out_indices, list) or isinstance(out_indices, tuple): self.out_indices = out_indices else: - raise TypeError('out_indices must be type of int, list or tuple') - - dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth) - ] # stochastic depth decay rule - self.blocks = nn.ModuleList([ - Block( - dim=embed_dim, - num_heads=num_heads, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - drop=dpr[i], - attn_drop=attn_drop_rate, - act_cfg=act_cfg, - norm_cfg=norm_cfg, - with_cp=with_cp) for i in range(depth) - ]) + raise TypeError("out_indices must be type of int, list or tuple") + + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule + self.blocks = nn.ModuleList( + [ + Block( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=dpr[i], + attn_drop=attn_drop_rate, + act_cfg=act_cfg, + norm_cfg=norm_cfg, + with_cp=with_cp, + ) + for i in range(depth) + ] + ) self.interpolate_mode = interpolate_mode self.final_norm = final_norm @@ -315,46 +302,47 @@ def init_weights(self, pretrained=None): if isinstance(pretrained, str): logger = get_root_logger() checkpoint = _load_checkpoint(pretrained, logger=logger) - if 'state_dict' in checkpoint: - state_dict = checkpoint['state_dict'] + if "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] else: state_dict = checkpoint - if 'pos_embed' in state_dict.keys(): - if self.pos_embed.shape != state_dict['pos_embed'].shape: - logger.info(msg=f'Resize the pos_embed shape from \ -{state_dict["pos_embed"].shape} to {self.pos_embed.shape}') + if "pos_embed" in state_dict.keys(): + if self.pos_embed.shape != state_dict["pos_embed"].shape: + logger.info( + msg=f'Resize the pos_embed shape from \ +{state_dict["pos_embed"].shape} to {self.pos_embed.shape}' + ) h, w = self.img_size - pos_size = int( - math.sqrt(state_dict['pos_embed'].shape[1] - 1)) - state_dict['pos_embed'] = self.resize_pos_embed( - state_dict['pos_embed'], (h, w), (pos_size, pos_size), - self.patch_size, self.interpolate_mode) + pos_size = int(math.sqrt(state_dict["pos_embed"].shape[1] - 1)) + state_dict["pos_embed"] = self.resize_pos_embed( + state_dict["pos_embed"], (h, w), (pos_size, pos_size), self.patch_size, self.interpolate_mode + ) self.load_state_dict(state_dict, False) elif pretrained is None: # We only implement the 'jax_impl' initialization implemented at # https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py#L353 # noqa: E501 - trunc_normal_(self.pos_embed, std=.02) - trunc_normal_(self.cls_token, std=.02) + trunc_normal_(self.pos_embed, std=0.02) + trunc_normal_(self.cls_token, std=0.02) for n, m in self.named_modules(): if isinstance(m, Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if m.bias is not None: - if 'mlp' in n: + if "mlp" in n: normal_init(m.bias, std=1e-6) else: constant_init(m.bias, 0) elif isinstance(m, Conv2d): - kaiming_init(m.weight, mode='fan_in') + kaiming_init(m.weight, mode="fan_in") if m.bias is not None: constant_init(m.bias, 0) elif isinstance(m, (_BatchNorm, nn.GroupNorm, nn.LayerNorm)): constant_init(m.bias, 0) constant_init(m.weight, 1.0) else: - raise TypeError('pretrained must be a str or None') + raise TypeError("pretrained must be a str or None") def _pos_embeding(self, img, patched_img, pos_embed): """Positiong embeding method. @@ -371,21 +359,19 @@ def _pos_embeding(self, img, patched_img, pos_embed): Return: torch.Tensor: The pos encoded image feature. """ - assert patched_img.ndim == 3 and pos_embed.ndim == 3, \ - 'the shapes of patched_img and pos_embed must be [B, L, C]' + assert ( + patched_img.ndim == 3 and pos_embed.ndim == 3 + ), "the shapes of patched_img and pos_embed must be [B, L, C]" x_len, pos_len = patched_img.shape[1], pos_embed.shape[1] if x_len != pos_len: - if pos_len == (self.img_size[0] // self.patch_size) * ( - self.img_size[1] // self.patch_size) + 1: + if pos_len == (self.img_size[0] // self.patch_size) * (self.img_size[1] // self.patch_size) + 1: pos_h = self.img_size[0] // self.patch_size pos_w = self.img_size[1] // self.patch_size else: - raise ValueError( - 'Unexpected shape of pos_embed, got {}.'.format( - pos_embed.shape)) - pos_embed = self.resize_pos_embed(pos_embed, img.shape[2:], - (pos_h, pos_w), self.patch_size, - self.interpolate_mode) + raise ValueError("Unexpected shape of pos_embed, got {}.".format(pos_embed.shape)) + pos_embed = self.resize_pos_embed( + pos_embed, img.shape[2:], (pos_h, pos_w), self.patch_size, self.interpolate_mode + ) return self.pos_drop(patched_img + pos_embed) @staticmethod @@ -401,18 +387,15 @@ def resize_pos_embed(pos_embed, input_shpae, pos_shape, patch_size, mode): Return: torch.Tensor: The resized pos_embed of shape [B, L_new, C] """ - assert pos_embed.ndim == 3, 'shape of pos_embed must be [B, L, C]' + assert pos_embed.ndim == 3, "shape of pos_embed must be [B, L, C]" input_h, input_w = input_shpae pos_h, pos_w = pos_shape cls_token_weight = pos_embed[:, 0] - pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w):] - pos_embed_weight = pos_embed_weight.reshape( - 1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2) + pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w) :] + pos_embed_weight = pos_embed_weight.reshape(1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2) pos_embed_weight = F.interpolate( - pos_embed_weight, - size=[input_h // patch_size, input_w // patch_size], - align_corners=False, - mode=mode) + pos_embed_weight, size=[input_h // patch_size, input_w // patch_size], align_corners=False, mode=mode + ) cls_token_weight = cls_token_weight.unsqueeze(1) pos_embed_weight = torch.flatten(pos_embed_weight, 2).transpose(1, 2) pos_embed = torch.cat((cls_token_weight, pos_embed_weight), dim=1) @@ -444,9 +427,9 @@ def forward(self, inputs): else: out = x B, _, C = out.shape - out = out.reshape(B, inputs.shape[2] // self.patch_size, - inputs.shape[3] // self.patch_size, - C).permute(0, 3, 1, 2) + out = out.reshape( + B, inputs.shape[2] // self.patch_size, inputs.shape[3] // self.patch_size, C + ).permute(0, 3, 1, 2) outs.append(out) return tuple(outs) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/builder.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/builder.py index 1cde93b6..5eb6a081 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/builder.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/builder.py @@ -1,8 +1,8 @@ import sys import warnings -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import MODELS as MMCV_MODELS -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import MODELS as MMCV_MODELS +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import Registry MODELS = Registry("models", sys.modules[__name__], parent=MMCV_MODELS) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ann_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ann_head.py index 3734959c..de3a6cdc 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ann_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ann_head.py @@ -1,6 +1,6 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule from ..builder import HEADS from ..utils import SelfAttentionBlock as _SelfAttentionBlock @@ -16,8 +16,7 @@ class PPMConcat(nn.ModuleList): """ def __init__(self, pool_scales=(1, 3, 6, 8)): - super(PPMConcat, self).__init__( - [nn.AdaptiveAvgPool2d(pool_scale) for pool_scale in pool_scales]) + super(PPMConcat, self).__init__([nn.AdaptiveAvgPool2d(pool_scale) for pool_scale in pool_scales]) def forward(self, feats): """Forward function.""" @@ -49,9 +48,19 @@ class SelfAttentionBlock(_SelfAttentionBlock): act_cfg (dict|None): Config of activation layers. """ - def __init__(self, low_in_channels, high_in_channels, channels, - out_channels, share_key_query, query_scale, key_pool_scales, - conv_cfg, norm_cfg, act_cfg): + def __init__( + self, + low_in_channels, + high_in_channels, + channels, + out_channels, + share_key_query, + query_scale, + key_pool_scales, + conv_cfg, + norm_cfg, + act_cfg, + ): key_psp = PPMConcat(key_pool_scales) if query_scale > 1: query_downsample = nn.MaxPool2d(kernel_size=query_scale) @@ -73,7 +82,8 @@ def __init__(self, low_in_channels, high_in_channels, channels, with_out=True, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) class AFNB(nn.Module): @@ -96,9 +106,18 @@ class AFNB(nn.Module): act_cfg (dict|None): Config of activation layers. """ - def __init__(self, low_in_channels, high_in_channels, channels, - out_channels, query_scales, key_pool_scales, conv_cfg, - norm_cfg, act_cfg): + def __init__( + self, + low_in_channels, + high_in_channels, + channels, + out_channels, + query_scales, + key_pool_scales, + conv_cfg, + norm_cfg, + act_cfg, + ): super(AFNB, self).__init__() self.stages = nn.ModuleList() for query_scale in query_scales: @@ -113,14 +132,12 @@ def __init__(self, low_in_channels, high_in_channels, channels, key_pool_scales=key_pool_scales, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg)) + act_cfg=act_cfg, + ) + ) self.bottleneck = ConvModule( - out_channels + high_in_channels, - out_channels, - 1, - conv_cfg=conv_cfg, - norm_cfg=norm_cfg, - act_cfg=None) + out_channels + high_in_channels, out_channels, 1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=None + ) def forward(self, low_feats, high_feats): """Forward function.""" @@ -147,8 +164,9 @@ class APNB(nn.Module): act_cfg (dict|None): Config of activation layers. """ - def __init__(self, in_channels, channels, out_channels, query_scales, - key_pool_scales, conv_cfg, norm_cfg, act_cfg): + def __init__( + self, in_channels, channels, out_channels, query_scales, key_pool_scales, conv_cfg, norm_cfg, act_cfg + ): super(APNB, self).__init__() self.stages = nn.ModuleList() for query_scale in query_scales: @@ -163,14 +181,12 @@ def __init__(self, in_channels, channels, out_channels, query_scales, key_pool_scales=key_pool_scales, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg)) + act_cfg=act_cfg, + ) + ) self.bottleneck = ConvModule( - 2 * in_channels, - out_channels, - 1, - conv_cfg=conv_cfg, - norm_cfg=norm_cfg, - act_cfg=act_cfg) + 2 * in_channels, out_channels, 1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg + ) def forward(self, feats): """Forward function.""" @@ -195,13 +211,8 @@ class ANNHead(BaseDecodeHead): Default: (1, 3, 6, 8). """ - def __init__(self, - project_channels, - query_scales=(1, ), - key_pool_scales=(1, 3, 6, 8), - **kwargs): - super(ANNHead, self).__init__( - input_transform='multiple_select', **kwargs) + def __init__(self, project_channels, query_scales=(1,), key_pool_scales=(1, 3, 6, 8), **kwargs): + super(ANNHead, self).__init__(input_transform="multiple_select", **kwargs) assert len(self.in_channels) == 2 low_in_channels, high_in_channels = self.in_channels self.project_channels = project_channels @@ -214,7 +225,8 @@ def __init__(self, key_pool_scales=key_pool_scales, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.bottleneck = ConvModule( high_in_channels, self.channels, @@ -222,7 +234,8 @@ def __init__(self, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.context = APNB( in_channels=self.channels, out_channels=self.channels, @@ -231,7 +244,8 @@ def __init__(self, key_pool_scales=key_pool_scales, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) def forward(self, inputs): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/apc_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/apc_head.py index 5f16be76..6a820b57 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/apc_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/apc_head.py @@ -1,9 +1,9 @@ import torch import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import HEADS from .decode_head import BaseDecodeHead @@ -22,8 +22,7 @@ class ACM(nn.Module): act_cfg (dict): Config of activation layers. """ - def __init__(self, pool_scale, fusion, in_channels, channels, conv_cfg, - norm_cfg, act_cfg): + def __init__(self, pool_scale, fusion, in_channels, channels, conv_cfg, norm_cfg, act_cfg): super(ACM, self).__init__() self.pool_scale = pool_scale self.fusion = fusion @@ -33,47 +32,27 @@ def __init__(self, pool_scale, fusion, in_channels, channels, conv_cfg, self.norm_cfg = norm_cfg self.act_cfg = act_cfg self.pooled_redu_conv = ConvModule( - self.in_channels, - self.channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + self.in_channels, self.channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) self.input_redu_conv = ConvModule( - self.in_channels, - self.channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + self.in_channels, self.channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) self.global_info = ConvModule( - self.channels, - self.channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + self.channels, self.channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) self.gla = nn.Conv2d(self.channels, self.pool_scale**2, 1, 1, 0) self.residual_conv = ConvModule( - self.channels, - self.channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + self.channels, self.channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) if self.fusion: self.fusion_conv = ConvModule( - self.channels, - self.channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + self.channels, self.channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) def forward(self, x): """Forward function.""" @@ -84,13 +63,13 @@ def forward(self, x): pooled_x = self.pooled_redu_conv(pooled_x) batch_size = x.size(0) # [batch_size, pool_scale * pool_scale, channels] - pooled_x = pooled_x.view(batch_size, self.channels, - -1).permute(0, 2, 1).contiguous() + pooled_x = pooled_x.view(batch_size, self.channels, -1).permute(0, 2, 1).contiguous() # [batch_size, h * w, pool_scale * pool_scale] - affinity_matrix = self.gla(x + resize( - self.global_info(F.adaptive_avg_pool2d(x, 1)), size=x.shape[2:]) - ).permute(0, 2, 3, 1).reshape( - batch_size, -1, self.pool_scale**2) + affinity_matrix = ( + self.gla(x + resize(self.global_info(F.adaptive_avg_pool2d(x, 1)), size=x.shape[2:])) + .permute(0, 2, 3, 1) + .reshape(batch_size, -1, self.pool_scale**2) + ) affinity_matrix = F.sigmoid(affinity_matrix) # [batch_size, h * w, channels] z_out = torch.matmul(affinity_matrix, pooled_x) @@ -129,13 +108,16 @@ def __init__(self, pool_scales=(1, 2, 3, 6), fusion=True, **kwargs): acm_modules = [] for pool_scale in self.pool_scales: acm_modules.append( - ACM(pool_scale, + ACM( + pool_scale, self.fusion, self.in_channels, self.channels, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg)) + act_cfg=self.act_cfg, + ) + ) self.acm_modules = nn.ModuleList(acm_modules) self.bottleneck = ConvModule( self.in_channels + len(pool_scales) * self.channels, @@ -144,7 +126,8 @@ def __init__(self, pool_scales=(1, 2, 3, 6), fusion=True, **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) def forward(self, inputs): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/aspp_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/aspp_head.py index ce144151..4d719f8b 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/aspp_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/aspp_head.py @@ -1,8 +1,8 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import HEADS from .decode_head import BaseDecodeHead @@ -19,8 +19,7 @@ class ASPPModule(nn.ModuleList): act_cfg (dict): Config of activation layers. """ - def __init__(self, dilations, in_channels, channels, conv_cfg, norm_cfg, - act_cfg): + def __init__(self, dilations, in_channels, channels, conv_cfg, norm_cfg, act_cfg): super(ASPPModule, self).__init__() self.dilations = dilations self.in_channels = in_channels @@ -38,7 +37,9 @@ def __init__(self, dilations, in_channels, channels, conv_cfg, norm_cfg, padding=0 if dilation == 1 else dilation, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg)) + act_cfg=self.act_cfg, + ) + ) def forward(self, x): """Forward function.""" @@ -73,14 +74,17 @@ def __init__(self, dilations=(1, 6, 12, 18), **kwargs): 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg)) + act_cfg=self.act_cfg, + ), + ) self.aspp_modules = ASPPModule( dilations, self.in_channels, self.channels, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.bottleneck = ConvModule( (len(dilations) + 1) * self.channels, self.channels, @@ -88,18 +92,13 @@ def __init__(self, dilations=(1, 6, 12, 18), **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) def forward(self, inputs): """Forward function.""" x = self._transform_inputs(inputs) - aspp_outs = [ - resize( - self.image_pool(x), - size=x.size()[2:], - mode='bilinear', - align_corners=self.align_corners) - ] + aspp_outs = [resize(self.image_pool(x), size=x.size()[2:], mode="bilinear", align_corners=self.align_corners)] aspp_outs.extend(self.aspp_modules(x)) aspp_outs = torch.cat(aspp_outs, dim=1) output = self.bottleneck(aspp_outs) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/cc_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/cc_head.py index a912c2a9..dd165cc3 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/cc_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/cc_head.py @@ -4,7 +4,7 @@ from .fcn_head import FCNHead try: - from comfy_controlnet_preprocessors.uniformer.mmcv.ops import CrissCrossAttention + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.ops import CrissCrossAttention except ModuleNotFoundError: CrissCrossAttention = None @@ -23,8 +23,7 @@ class CCHead(FCNHead): def __init__(self, recurrence=2, **kwargs): if CrissCrossAttention is None: - raise RuntimeError('Please install mmcv-full for ' - 'CrissCrossAttention ops') + raise RuntimeError("Please install mmcv-full for " "CrissCrossAttention ops") super(CCHead, self).__init__(num_convs=2, **kwargs) self.recurrence = recurrence self.cca = CrissCrossAttention(self.channels) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/da_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/da_head.py index 946c308e..1f00b5f9 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/da_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/da_head.py @@ -1,9 +1,9 @@ import torch import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, Scale +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, Scale from torch import nn -from comfy_controlnet_preprocessors.uniformer.mmseg.core import add_prefix +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.core import add_prefix from ..builder import HEADS from ..utils import SelfAttentionBlock as _SelfAttentionBlock from .decode_head import BaseDecodeHead @@ -34,7 +34,8 @@ def __init__(self, in_channels, channels): with_out=False, conv_cfg=None, norm_cfg=None, - act_cfg=None) + act_cfg=None, + ) self.gamma = Scale(0) @@ -59,8 +60,7 @@ def forward(self, x): proj_query = x.view(batch_size, channels, -1) proj_key = x.view(batch_size, channels, -1).permute(0, 2, 1) energy = torch.bmm(proj_query, proj_key) - energy_new = torch.max( - energy, -1, keepdim=True)[0].expand_as(energy) - energy + energy_new = torch.max(energy, -1, keepdim=True)[0].expand_as(energy) - energy attention = F.softmax(energy_new, dim=-1) proj_value = x.view(batch_size, channels, -1) @@ -92,7 +92,8 @@ def __init__(self, pam_channels, **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.pam = PAM(self.channels, pam_channels) self.pam_out_conv = ConvModule( self.channels, @@ -101,9 +102,9 @@ def __init__(self, pam_channels, **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) - self.pam_conv_seg = nn.Conv2d( - self.channels, self.num_classes, kernel_size=1) + act_cfg=self.act_cfg, + ) + self.pam_conv_seg = nn.Conv2d(self.channels, self.num_classes, kernel_size=1) self.cam_in_conv = ConvModule( self.in_channels, @@ -112,7 +113,8 @@ def __init__(self, pam_channels, **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.cam = CAM() self.cam_out_conv = ConvModule( self.channels, @@ -121,9 +123,9 @@ def __init__(self, pam_channels, **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) - self.cam_conv_seg = nn.Conv2d( - self.channels, self.num_classes, kernel_size=1) + act_cfg=self.act_cfg, + ) + self.cam_conv_seg = nn.Conv2d(self.channels, self.num_classes, kernel_size=1) def pam_cls_seg(self, feat): """PAM feature classification.""" @@ -165,14 +167,7 @@ def losses(self, seg_logit, seg_label): """Compute ``pam_cam``, ``pam``, ``cam`` loss.""" pam_cam_seg_logit, pam_seg_logit, cam_seg_logit = seg_logit loss = dict() - loss.update( - add_prefix( - super(DAHead, self).losses(pam_cam_seg_logit, seg_label), - 'pam_cam')) - loss.update( - add_prefix( - super(DAHead, self).losses(pam_seg_logit, seg_label), 'pam')) - loss.update( - add_prefix( - super(DAHead, self).losses(cam_seg_logit, seg_label), 'cam')) + loss.update(add_prefix(super(DAHead, self).losses(pam_cam_seg_logit, seg_label), "pam_cam")) + loss.update(add_prefix(super(DAHead, self).losses(pam_seg_logit, seg_label), "pam")) + loss.update(add_prefix(super(DAHead, self).losses(cam_seg_logit, seg_label), "cam")) return loss diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/decode_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/decode_head.py index 047da4fb..9756106c 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/decode_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/decode_head.py @@ -2,11 +2,11 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import normal_init -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import auto_fp16, force_fp32 +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import normal_init +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import auto_fp16, force_fp32 -from comfy_controlnet_preprocessors.uniformer.mmseg.core import build_pixel_sampler -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.core import build_pixel_sampler +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import build_loss from ..losses import accuracy @@ -43,24 +43,23 @@ class BaseDecodeHead(nn.Module, metaclass=ABCMeta): Default: False. """ - def __init__(self, - in_channels, - channels, - *, - num_classes, - dropout_ratio=0.1, - conv_cfg=None, - norm_cfg=None, - act_cfg=dict(type='ReLU'), - in_index=-1, - input_transform=None, - loss_decode=dict( - type='CrossEntropyLoss', - use_sigmoid=False, - loss_weight=1.0), - ignore_index=255, - sampler=None, - align_corners=False): + def __init__( + self, + in_channels, + channels, + *, + num_classes, + dropout_ratio=0.1, + conv_cfg=None, + norm_cfg=None, + act_cfg=dict(type="ReLU"), + in_index=-1, + input_transform=None, + loss_decode=dict(type="CrossEntropyLoss", use_sigmoid=False, loss_weight=1.0), + ignore_index=255, + sampler=None, + align_corners=False, + ): super(BaseDecodeHead, self).__init__() self._init_inputs(in_channels, in_index, input_transform) self.channels = channels @@ -87,9 +86,11 @@ def __init__(self, def extra_repr(self): """Extra repr.""" - s = f'input_transform={self.input_transform}, ' \ - f'ignore_index={self.ignore_index}, ' \ - f'align_corners={self.align_corners}' + s = ( + f"input_transform={self.input_transform}, " + f"ignore_index={self.ignore_index}, " + f"align_corners={self.align_corners}" + ) return s def _init_inputs(self, in_channels, in_index, input_transform): @@ -114,14 +115,14 @@ def _init_inputs(self, in_channels, in_index, input_transform): """ if input_transform is not None: - assert input_transform in ['resize_concat', 'multiple_select'] + assert input_transform in ["resize_concat", "multiple_select"] self.input_transform = input_transform self.in_index = in_index if input_transform is not None: assert isinstance(in_channels, (list, tuple)) assert isinstance(in_index, (list, tuple)) assert len(in_channels) == len(in_index) - if input_transform == 'resize_concat': + if input_transform == "resize_concat": self.in_channels = sum(in_channels) else: self.in_channels = in_channels @@ -144,17 +145,14 @@ def _transform_inputs(self, inputs): Tensor: The transformed inputs """ - if self.input_transform == 'resize_concat': + if self.input_transform == "resize_concat": inputs = [inputs[i] for i in self.in_index] upsampled_inputs = [ - resize( - input=x, - size=inputs[0].shape[2:], - mode='bilinear', - align_corners=self.align_corners) for x in inputs + resize(input=x, size=inputs[0].shape[2:], mode="bilinear", align_corners=self.align_corners) + for x in inputs ] inputs = torch.cat(upsampled_inputs, dim=1) - elif self.input_transform == 'multiple_select': + elif self.input_transform == "multiple_select": inputs = [inputs[i] for i in self.in_index] else: inputs = inputs[self.in_index] @@ -211,24 +209,18 @@ def cls_seg(self, feat): output = self.conv_seg(feat) return output - @force_fp32(apply_to=('seg_logit', )) + @force_fp32(apply_to=("seg_logit",)) def losses(self, seg_logit, seg_label): """Compute segmentation loss.""" loss = dict() seg_logit = resize( - input=seg_logit, - size=seg_label.shape[2:], - mode='bilinear', - align_corners=self.align_corners) + input=seg_logit, size=seg_label.shape[2:], mode="bilinear", align_corners=self.align_corners + ) if self.sampler is not None: seg_weight = self.sampler.sample(seg_logit, seg_label) else: seg_weight = None seg_label = seg_label.squeeze(1) - loss['loss_seg'] = self.loss_decode( - seg_logit, - seg_label, - weight=seg_weight, - ignore_index=self.ignore_index) - loss['acc_seg'] = accuracy(seg_logit, seg_label) + loss["loss_seg"] = self.loss_decode(seg_logit, seg_label, weight=seg_weight, ignore_index=self.ignore_index) + loss["acc_seg"] = accuracy(seg_logit, seg_label) return loss diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/dm_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/dm_head.py index 3ccacf62..ff5162c4 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/dm_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/dm_head.py @@ -1,7 +1,11 @@ import torch import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, build_activation_layer, build_norm_layer +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ( + ConvModule, + build_activation_layer, + build_norm_layer, +) from ..builder import HEADS from .decode_head import BaseDecodeHead @@ -21,8 +25,7 @@ class DCM(nn.Module): act_cfg (dict): Config of activation layers. """ - def __init__(self, filter_size, fusion, in_channels, channels, conv_cfg, - norm_cfg, act_cfg): + def __init__(self, filter_size, fusion, in_channels, channels, conv_cfg, norm_cfg, act_cfg): super(DCM, self).__init__() self.filter_size = filter_size self.fusion = fusion @@ -31,16 +34,11 @@ def __init__(self, filter_size, fusion, in_channels, channels, conv_cfg, self.conv_cfg = conv_cfg self.norm_cfg = norm_cfg self.act_cfg = act_cfg - self.filter_gen_conv = nn.Conv2d(self.in_channels, self.channels, 1, 1, - 0) + self.filter_gen_conv = nn.Conv2d(self.in_channels, self.channels, 1, 1, 0) self.input_redu_conv = ConvModule( - self.in_channels, - self.channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + self.in_channels, self.channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) if self.norm_cfg is not None: self.norm = build_norm_layer(self.norm_cfg, self.channels)[1] @@ -50,30 +48,24 @@ def __init__(self, filter_size, fusion, in_channels, channels, conv_cfg, if self.fusion: self.fusion_conv = ConvModule( - self.channels, - self.channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + self.channels, self.channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) def forward(self, x): """Forward function.""" - generated_filter = self.filter_gen_conv( - F.adaptive_avg_pool2d(x, self.filter_size)) + generated_filter = self.filter_gen_conv(F.adaptive_avg_pool2d(x, self.filter_size)) x = self.input_redu_conv(x) b, c, h, w = x.shape # [1, b * c, h, w], c = self.channels x = x.view(1, b * c, h, w) # [b * c, 1, filter_size, filter_size] - generated_filter = generated_filter.view(b * c, 1, self.filter_size, - self.filter_size) + generated_filter = generated_filter.view(b * c, 1, self.filter_size, self.filter_size) pad = (self.filter_size - 1) // 2 if (self.filter_size - 1) % 2 == 0: p2d = (pad, pad, pad, pad) else: p2d = (pad + 1, pad, pad + 1, pad) - x = F.pad(input=x, pad=p2d, mode='constant', value=0) + x = F.pad(input=x, pad=p2d, mode="constant", value=0) # [1, b * c, h, w] output = F.conv2d(input=x, weight=generated_filter, groups=b * c) # [b, c, h, w] @@ -111,13 +103,16 @@ def __init__(self, filter_sizes=(1, 3, 5, 7), fusion=False, **kwargs): dcm_modules = [] for filter_size in self.filter_sizes: dcm_modules.append( - DCM(filter_size, + DCM( + filter_size, self.fusion, self.in_channels, self.channels, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg)) + act_cfg=self.act_cfg, + ) + ) self.dcm_modules = nn.ModuleList(dcm_modules) self.bottleneck = ConvModule( self.in_channels + len(filter_sizes) * self.channels, @@ -126,7 +121,8 @@ def __init__(self, filter_sizes=(1, 3, 5, 7), fusion=False, **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) def forward(self, inputs): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/dnl_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/dnl_head.py index 2ce60941..df294abe 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/dnl_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/dnl_head.py @@ -1,5 +1,5 @@ import torch -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import NonLocal2d +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import NonLocal2d from torch import nn from ..builder import HEADS @@ -25,7 +25,7 @@ def embedded_gaussian(self, theta_x, phi_x): pairwise_weight = torch.matmul(theta_x, phi_x) if self.use_scale: # theta_x.shape[-1] is `self.inter_channels` - pairwise_weight /= theta_x.shape[-1]**0.5 + pairwise_weight /= theta_x.shape[-1] ** 0.5 pairwise_weight /= self.temperature pairwise_weight = pairwise_weight.softmax(dim=-1) return pairwise_weight @@ -39,14 +39,14 @@ def forward(self, x): g_x = g_x.permute(0, 2, 1) # theta_x: [N, HxW, C], phi_x: [N, C, HxW] - if self.mode == 'gaussian': + if self.mode == "gaussian": theta_x = x.view(n, self.in_channels, -1) theta_x = theta_x.permute(0, 2, 1) if self.sub_sample: phi_x = self.phi(x).view(n, self.in_channels, -1) else: phi_x = x.view(n, self.in_channels, -1) - elif self.mode == 'concatenation': + elif self.mode == "concatenation": theta_x = self.theta(x).view(n, self.inter_channels, -1, 1) phi_x = self.phi(x).view(n, self.inter_channels, 1, -1) else: @@ -65,8 +65,7 @@ def forward(self, x): # y: [N, HxW, C] y = torch.matmul(pairwise_weight, g_x) # y: [N, C, H, W] - y = y.permute(0, 2, 1).contiguous().reshape(n, self.inter_channels, - *x.size()[2:]) + y = y.permute(0, 2, 1).contiguous().reshape(n, self.inter_channels, *x.size()[2:]) # unary_mask: [N, 1, HxW] unary_mask = self.conv_mask(x) @@ -75,8 +74,7 @@ def forward(self, x): # unary_x: [N, 1, C] unary_x = torch.matmul(unary_mask, g_x) # unary_x: [N, C, 1, 1] - unary_x = unary_x.permute(0, 2, 1).contiguous().reshape( - n, self.inter_channels, 1, 1) + unary_x = unary_x.permute(0, 2, 1).contiguous().reshape(n, self.inter_channels, 1, 1) output = x + self.conv_out(y + unary_x) @@ -99,12 +97,7 @@ class DNLHead(FCNHead): temperature (float): Temperature to adjust attention. Default: 0.05 """ - def __init__(self, - reduction=2, - use_scale=True, - mode='embedded_gaussian', - temperature=0.05, - **kwargs): + def __init__(self, reduction=2, use_scale=True, mode="embedded_gaussian", temperature=0.05, **kwargs): super(DNLHead, self).__init__(num_convs=2, **kwargs) self.reduction = reduction self.use_scale = use_scale @@ -117,7 +110,8 @@ def __init__(self, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, mode=self.mode, - temperature=self.temperature) + temperature=self.temperature, + ) def forward(self, inputs): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ema_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ema_head.py index c1338527..2e3c6f9c 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ema_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ema_head.py @@ -4,7 +4,7 @@ import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule from ..builder import HEADS from .decode_head import BaseDecodeHead @@ -30,16 +30,16 @@ class EMAModule(nn.Module): def __init__(self, channels, num_bases, num_stages, momentum): super(EMAModule, self).__init__() - assert num_stages >= 1, 'num_stages must be at least 1!' + assert num_stages >= 1, "num_stages must be at least 1!" self.num_bases = num_bases self.num_stages = num_stages self.momentum = momentum bases = torch.zeros(1, channels, self.num_bases) - bases.normal_(0, math.sqrt(2. / self.num_bases)) + bases.normal_(0, math.sqrt(2.0 / self.num_bases)) # [1, channels, num_bases] bases = F.normalize(bases, dim=1, p=2) - self.register_buffer('bases', bases) + self.register_buffer("bases", bases) def forward(self, feats): """Forward function.""" @@ -52,16 +52,16 @@ def forward(self, feats): with torch.no_grad(): for i in range(self.num_stages): # [batch_size, height*width, num_bases] - attention = torch.einsum('bcn,bck->bnk', feats, bases) + attention = torch.einsum("bcn,bck->bnk", feats, bases) attention = F.softmax(attention, dim=2) # l1 norm attention_normed = F.normalize(attention, dim=1, p=1) # [batch_size, channels, num_bases] - bases = torch.einsum('bcn,bnk->bck', feats, attention_normed) + bases = torch.einsum("bcn,bnk->bck", feats, attention_normed) # l2 norm bases = F.normalize(bases, dim=1, p=2) - feats_recon = torch.einsum('bck,bnk->bcn', bases, attention) + feats_recon = torch.einsum("bck,bnk->bcn", bases, attention) feats_recon = feats_recon.view(batch_size, channels, height, width) if self.training: @@ -69,8 +69,7 @@ def forward(self, feats): bases = reduce_mean(bases) # l2 norm bases = F.normalize(bases, dim=1, p=2) - self.bases = (1 - - self.momentum) * self.bases + self.momentum * bases + self.bases = (1 - self.momentum) * self.bases + self.momentum * bases return feats_recon @@ -91,21 +90,14 @@ class EMAHead(BaseDecodeHead): momentum (float): Momentum to update the base. Default: 0.1. """ - def __init__(self, - ema_channels, - num_bases, - num_stages, - concat_input=True, - momentum=0.1, - **kwargs): + def __init__(self, ema_channels, num_bases, num_stages, concat_input=True, momentum=0.1, **kwargs): super(EMAHead, self).__init__(**kwargs) self.ema_channels = ema_channels self.num_bases = num_bases self.num_stages = num_stages self.concat_input = concat_input self.momentum = momentum - self.ema_module = EMAModule(self.ema_channels, self.num_bases, - self.num_stages, self.momentum) + self.ema_module = EMAModule(self.ema_channels, self.num_bases, self.num_stages, self.momentum) self.ema_in_conv = ConvModule( self.in_channels, @@ -114,25 +106,18 @@ def __init__(self, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) # project (0, inf) -> (-inf, inf) self.ema_mid_conv = ConvModule( - self.ema_channels, - self.ema_channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=None, - act_cfg=None) + self.ema_channels, self.ema_channels, 1, conv_cfg=self.conv_cfg, norm_cfg=None, act_cfg=None + ) for param in self.ema_mid_conv.parameters(): param.requires_grad = False self.ema_out_conv = ConvModule( - self.ema_channels, - self.ema_channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=None) + self.ema_channels, self.ema_channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=None + ) self.bottleneck = ConvModule( self.ema_channels, self.channels, @@ -140,7 +125,8 @@ def __init__(self, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) if self.concat_input: self.conv_cat = ConvModule( self.in_channels + self.channels, @@ -149,7 +135,8 @@ def __init__(self, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) def forward(self, inputs): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/enc_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/enc_head.py index eb2cd82c..abdcc481 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/enc_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/enc_head.py @@ -1,9 +1,9 @@ import torch import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, build_norm_layer +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, build_norm_layer -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import Encoding, resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import Encoding, resize from ..builder import HEADS, build_loss from .decode_head import BaseDecodeHead @@ -22,30 +22,25 @@ class EncModule(nn.Module): def __init__(self, in_channels, num_codes, conv_cfg, norm_cfg, act_cfg): super(EncModule, self).__init__() self.encoding_project = ConvModule( - in_channels, - in_channels, - 1, - conv_cfg=conv_cfg, - norm_cfg=norm_cfg, - act_cfg=act_cfg) + in_channels, in_channels, 1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg + ) # TODO: resolve this hack # change to 1d if norm_cfg is not None: encoding_norm_cfg = norm_cfg.copy() - if encoding_norm_cfg['type'] in ['BN', 'IN']: - encoding_norm_cfg['type'] += '1d' + if encoding_norm_cfg["type"] in ["BN", "IN"]: + encoding_norm_cfg["type"] += "1d" else: - encoding_norm_cfg['type'] = encoding_norm_cfg['type'].replace( - '2d', '1d') + encoding_norm_cfg["type"] = encoding_norm_cfg["type"].replace("2d", "1d") else: # fallback to BN1d - encoding_norm_cfg = dict(type='BN1d') + encoding_norm_cfg = dict(type="BN1d") self.encoding = nn.Sequential( Encoding(channels=in_channels, num_codes=num_codes), build_norm_layer(encoding_norm_cfg, num_codes)[1], - nn.ReLU(inplace=True)) - self.fc = nn.Sequential( - nn.Linear(in_channels, in_channels), nn.Sigmoid()) + nn.ReLU(inplace=True), + ) + self.fc = nn.Sequential(nn.Linear(in_channels, in_channels), nn.Sigmoid()) def forward(self, x): """Forward function.""" @@ -75,17 +70,15 @@ class EncHead(BaseDecodeHead): Default: dict(type='CrossEntropyLoss', use_sigmoid=True). """ - def __init__(self, - num_codes=32, - use_se_loss=True, - add_lateral=False, - loss_se_decode=dict( - type='CrossEntropyLoss', - use_sigmoid=True, - loss_weight=0.2), - **kwargs): - super(EncHead, self).__init__( - input_transform='multiple_select', **kwargs) + def __init__( + self, + num_codes=32, + use_se_loss=True, + add_lateral=False, + loss_se_decode=dict(type="CrossEntropyLoss", use_sigmoid=True, loss_weight=0.2), + **kwargs, + ): + super(EncHead, self).__init__(input_transform="multiple_select", **kwargs) self.use_se_loss = use_se_loss self.add_lateral = add_lateral self.num_codes = num_codes @@ -96,7 +89,8 @@ def __init__(self, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) if add_lateral: self.lateral_convs = nn.ModuleList() for in_channels in self.in_channels[:-1]: # skip the last one @@ -107,7 +101,9 @@ def __init__(self, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg)) + act_cfg=self.act_cfg, + ) + ) self.fusion = ConvModule( len(self.in_channels) * self.channels, self.channels, @@ -115,13 +111,11 @@ def __init__(self, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.enc_module = EncModule( - self.channels, - num_codes=num_codes, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + self.channels, num_codes=num_codes, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) if self.use_se_loss: self.loss_se_decode = build_loss(loss_se_decode) self.se_layer = nn.Linear(self.channels, self.num_classes) @@ -132,11 +126,7 @@ def forward(self, inputs): feat = self.bottleneck(inputs[-1]) if self.add_lateral: laterals = [ - resize( - lateral_conv(inputs[i]), - size=feat.shape[2:], - mode='bilinear', - align_corners=self.align_corners) + resize(lateral_conv(inputs[i]), size=feat.shape[2:], mode="bilinear", align_corners=self.align_corners) for i, lateral_conv in enumerate(self.lateral_convs) ] feat = self.fusion(torch.cat([feat, *laterals], 1)) @@ -170,8 +160,7 @@ def _convert_to_onehot_labels(seg_label, num_classes): batch_size = seg_label.size(0) onehot_labels = seg_label.new_zeros((batch_size, num_classes)) for i in range(batch_size): - hist = seg_label[i].float().histc( - bins=num_classes, min=0, max=num_classes - 1) + hist = seg_label[i].float().histc(bins=num_classes, min=0, max=num_classes - 1) onehot_labels[i] = hist > 0 return onehot_labels @@ -180,8 +169,6 @@ def losses(self, seg_logit, seg_label): seg_logit, se_seg_logit = seg_logit loss = dict() loss.update(super(EncHead, self).losses(seg_logit, seg_label)) - se_loss = self.loss_se_decode( - se_seg_logit, - self._convert_to_onehot_labels(seg_label, self.num_classes)) - loss['loss_se'] = se_loss + se_loss = self.loss_se_decode(se_seg_logit, self._convert_to_onehot_labels(seg_label, self.num_classes)) + loss["loss_se"] = se_loss return loss diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/fcn_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/fcn_head.py index 6ef59644..2ae6d6f0 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/fcn_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/fcn_head.py @@ -1,6 +1,6 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule from ..builder import HEADS from .decode_head import BaseDecodeHead @@ -20,12 +20,7 @@ class FCNHead(BaseDecodeHead): dilation (int): The dilation rate for convs in the head. Default: 1. """ - def __init__(self, - num_convs=2, - kernel_size=3, - concat_input=True, - dilation=1, - **kwargs): + def __init__(self, num_convs=2, kernel_size=3, concat_input=True, dilation=1, **kwargs): assert num_convs >= 0 and dilation > 0 and isinstance(dilation, int) self.num_convs = num_convs self.concat_input = concat_input @@ -45,7 +40,9 @@ def __init__(self, dilation=dilation, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg)) + act_cfg=self.act_cfg, + ) + ) for i in range(num_convs - 1): convs.append( ConvModule( @@ -56,7 +53,9 @@ def __init__(self, dilation=dilation, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg)) + act_cfg=self.act_cfg, + ) + ) if num_convs == 0: self.convs = nn.Identity() else: @@ -69,7 +68,8 @@ def __init__(self, padding=kernel_size // 2, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) def forward(self, inputs): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/fpn_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/fpn_head.py index fd44f023..38d06bc1 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/fpn_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/fpn_head.py @@ -1,8 +1,8 @@ import numpy as np import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import HEADS from .decode_head import BaseDecodeHead @@ -21,17 +21,14 @@ class FPNHead(BaseDecodeHead): """ def __init__(self, feature_strides, **kwargs): - super(FPNHead, self).__init__( - input_transform='multiple_select', **kwargs) + super(FPNHead, self).__init__(input_transform="multiple_select", **kwargs) assert len(feature_strides) == len(self.in_channels) assert min(feature_strides) == feature_strides[0] self.feature_strides = feature_strides self.scale_heads = nn.ModuleList() for i in range(len(feature_strides)): - head_length = max( - 1, - int(np.log2(feature_strides[i]) - np.log2(feature_strides[0]))) + head_length = max(1, int(np.log2(feature_strides[i]) - np.log2(feature_strides[0]))) scale_head = [] for k in range(head_length): scale_head.append( @@ -42,27 +39,22 @@ def __init__(self, feature_strides, **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg)) + act_cfg=self.act_cfg, + ) + ) if feature_strides[i] != feature_strides[0]: - scale_head.append( - nn.Upsample( - scale_factor=2, - mode='bilinear', - align_corners=self.align_corners)) + scale_head.append(nn.Upsample(scale_factor=2, mode="bilinear", align_corners=self.align_corners)) self.scale_heads.append(nn.Sequential(*scale_head)) def forward(self, inputs): - x = self._transform_inputs(inputs) output = self.scale_heads[0](x[0]) for i in range(1, len(self.feature_strides)): # non inplace output = output + resize( - self.scale_heads[i](x[i]), - size=output.shape[2:], - mode='bilinear', - align_corners=self.align_corners) + self.scale_heads[i](x[i]), size=output.shape[2:], mode="bilinear", align_corners=self.align_corners + ) output = self.cls_seg(output) return output diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/gc_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/gc_head.py index 1e6ed9b4..e3b5e47d 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/gc_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/gc_head.py @@ -1,5 +1,5 @@ import torch -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ContextBlock +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ContextBlock from ..builder import HEADS from .fcn_head import FCNHead @@ -20,20 +20,14 @@ class GCHead(FCNHead): Options are 'channel_add', 'channel_mul'. Default: ('channel_add',) """ - def __init__(self, - ratio=1 / 4., - pooling_type='att', - fusion_types=('channel_add', ), - **kwargs): + def __init__(self, ratio=1 / 4.0, pooling_type="att", fusion_types=("channel_add",), **kwargs): super(GCHead, self).__init__(num_convs=2, **kwargs) self.ratio = ratio self.pooling_type = pooling_type self.fusion_types = fusion_types self.gc_block = ContextBlock( - in_channels=self.channels, - ratio=self.ratio, - pooling_type=self.pooling_type, - fusion_types=self.fusion_types) + in_channels=self.channels, ratio=self.ratio, pooling_type=self.pooling_type, fusion_types=self.fusion_types + ) def forward(self, inputs): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/lraspp_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/lraspp_head.py index 64cff660..639fd6bc 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/lraspp_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/lraspp_head.py @@ -1,9 +1,9 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv import is_tuple_of -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv import is_tuple_of +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import HEADS from .decode_head import BaseDecodeHead @@ -22,10 +22,12 @@ class LRASPPHead(BaseDecodeHead): def __init__(self, branch_channels=(32, 64), **kwargs): super(LRASPPHead, self).__init__(**kwargs) - if self.input_transform != 'multiple_select': - raise ValueError('in Lite R-ASPP (LRASPP) head, input_transform ' - f'must be \'multiple_select\'. But received ' - f'\'{self.input_transform}\'') + if self.input_transform != "multiple_select": + raise ValueError( + "in Lite R-ASPP (LRASPP) head, input_transform " + f"must be 'multiple_select'. But received " + f"'{self.input_transform}'" + ) assert is_tuple_of(branch_channels, int) assert len(branch_channels) == len(self.in_channels) - 1 self.branch_channels = branch_channels @@ -33,37 +35,28 @@ def __init__(self, branch_channels=(32, 64), **kwargs): self.convs = nn.Sequential() self.conv_ups = nn.Sequential() for i in range(len(branch_channels)): - self.convs.add_module( - f'conv{i}', - nn.Conv2d( - self.in_channels[i], branch_channels[i], 1, bias=False)) + self.convs.add_module(f"conv{i}", nn.Conv2d(self.in_channels[i], branch_channels[i], 1, bias=False)) self.conv_ups.add_module( - f'conv_up{i}', + f"conv_up{i}", ConvModule( self.channels + branch_channels[i], self.channels, 1, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg, - bias=False)) + bias=False, + ), + ) self.conv_up_input = nn.Conv2d(self.channels, self.channels, 1) self.aspp_conv = ConvModule( - self.in_channels[-1], - self.channels, - 1, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg, - bias=False) + self.in_channels[-1], self.channels, 1, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg, bias=False + ) self.image_pool = nn.Sequential( nn.AvgPool2d(kernel_size=49, stride=(16, 20)), - ConvModule( - self.in_channels[2], - self.channels, - 1, - act_cfg=dict(type='Sigmoid'), - bias=False)) + ConvModule(self.in_channels[2], self.channels, 1, act_cfg=dict(type="Sigmoid"), bias=False), + ) def forward(self, inputs): """Forward function.""" @@ -72,18 +65,12 @@ def forward(self, inputs): x = inputs[-1] x = self.aspp_conv(x) * resize( - self.image_pool(x), - size=x.size()[2:], - mode='bilinear', - align_corners=self.align_corners) + self.image_pool(x), size=x.size()[2:], mode="bilinear", align_corners=self.align_corners + ) x = self.conv_up_input(x) for i in range(len(self.branch_channels) - 1, -1, -1): - x = resize( - x, - size=inputs[i].size()[2:], - mode='bilinear', - align_corners=self.align_corners) + x = resize(x, size=inputs[i].size()[2:], mode="bilinear", align_corners=self.align_corners) x = torch.cat([x, self.convs[i](inputs[i])], 1) x = self.conv_ups[i](x) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/nl_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/nl_head.py index 3be36b21..ff8a5dce 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/nl_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/nl_head.py @@ -1,5 +1,5 @@ import torch -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import NonLocal2d +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import NonLocal2d from ..builder import HEADS from .fcn_head import FCNHead @@ -20,11 +20,7 @@ class NLHead(FCNHead): 'dot_product'. Default: 'embedded_gaussian.'. """ - def __init__(self, - reduction=2, - use_scale=True, - mode='embedded_gaussian', - **kwargs): + def __init__(self, reduction=2, use_scale=True, mode="embedded_gaussian", **kwargs): super(NLHead, self).__init__(num_convs=2, **kwargs) self.reduction = reduction self.use_scale = use_scale @@ -35,7 +31,8 @@ def __init__(self, use_scale=self.use_scale, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - mode=self.mode) + mode=self.mode, + ) def forward(self, inputs): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ocr_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ocr_head.py index 39f0821f..c9d9364c 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ocr_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/ocr_head.py @@ -1,9 +1,9 @@ import torch import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import HEADS from ..utils import SelfAttentionBlock as _SelfAttentionBlock from .cascade_decode_head import BaseCascadeDecodeHead @@ -39,8 +39,7 @@ def forward(self, feats, probs): class ObjectAttentionBlock(_SelfAttentionBlock): """Make a OCR used SelfAttentionBlock.""" - def __init__(self, in_channels, channels, scale, conv_cfg, norm_cfg, - act_cfg): + def __init__(self, in_channels, channels, scale, conv_cfg, norm_cfg, act_cfg): if scale > 1: query_downsample = nn.MaxPool2d(kernel_size=scale) else: @@ -61,19 +60,15 @@ def __init__(self, in_channels, channels, scale, conv_cfg, norm_cfg, with_out=True, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) self.bottleneck = ConvModule( - in_channels * 2, - in_channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + in_channels * 2, in_channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) def forward(self, query_feats, key_feats): """Forward function.""" - context = super(ObjectAttentionBlock, - self).forward(query_feats, key_feats) + context = super(ObjectAttentionBlock, self).forward(query_feats, key_feats) output = self.bottleneck(torch.cat([context, query_feats], dim=1)) if self.query_downsample is not None: output = resize(query_feats) @@ -104,7 +99,8 @@ def __init__(self, ocr_channels, scale=1, **kwargs): self.scale, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.spatial_gather_module = SpatialGatherModule(self.scale) self.bottleneck = ConvModule( @@ -114,7 +110,8 @@ def __init__(self, ocr_channels, scale=1, **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) def forward(self, inputs, prev_output): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/point_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/point_head.py index 8a32afdc..a61e4f84 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/point_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/point_head.py @@ -2,11 +2,11 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, normal_init -from comfy_controlnet_preprocessors.uniformer.mmcv.ops import point_sample +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, normal_init +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.ops import point_sample -from comfy_controlnet_preprocessors.uniformer.mmseg.models.builder import HEADS -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.models.builder import HEADS +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..losses import accuracy from .cascade_decode_head import BaseCascadeDecodeHead @@ -57,19 +57,18 @@ class PointHead(BaseCascadeDecodeHead): loss_weight=1.0). """ - def __init__(self, - num_fcs=3, - coarse_pred_each_layer=True, - conv_cfg=dict(type='Conv1d'), - norm_cfg=None, - act_cfg=dict(type='ReLU', inplace=False), - **kwargs): + def __init__( + self, + num_fcs=3, + coarse_pred_each_layer=True, + conv_cfg=dict(type="Conv1d"), + norm_cfg=None, + act_cfg=dict(type="ReLU", inplace=False), + **kwargs, + ): super(PointHead, self).__init__( - input_transform='multiple_select', - conv_cfg=conv_cfg, - norm_cfg=norm_cfg, - act_cfg=act_cfg, - **kwargs) + input_transform="multiple_select", conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg, **kwargs + ) self.num_fcs = num_fcs self.coarse_pred_each_layer = coarse_pred_each_layer @@ -86,20 +85,15 @@ def __init__(self, padding=0, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) self.fcs.append(fc) fc_in_channels = fc_channels - fc_in_channels += self.num_classes if self.coarse_pred_each_layer \ - else 0 - self.fc_seg = nn.Conv1d( - fc_in_channels, - self.num_classes, - kernel_size=1, - stride=1, - padding=0) + fc_in_channels += self.num_classes if self.coarse_pred_each_layer else 0 + self.fc_seg = nn.Conv1d(fc_in_channels, self.num_classes, kernel_size=1, stride=1, padding=0) if self.dropout_ratio > 0: self.dropout = nn.Dropout(self.dropout_ratio) - delattr(self, 'conv_seg') + delattr(self, "conv_seg") def init_weights(self): """Initialize weights of classification layer.""" @@ -133,10 +127,7 @@ def _get_fine_grained_point_feats(self, x, points): shape (batch_size, sum(channels of x), num_points). """ - fine_grained_feats_list = [ - point_sample(_, points, align_corners=self.align_corners) - for _ in x - ] + fine_grained_feats_list = [point_sample(_, points, align_corners=self.align_corners) for _ in x] if len(fine_grained_feats_list) > 1: fine_grained_feats = torch.cat(fine_grained_feats_list, dim=1) else: @@ -157,13 +148,11 @@ def _get_coarse_point_feats(self, prev_output, points): num_classes, num_points). """ - coarse_feats = point_sample( - prev_output, points, align_corners=self.align_corners) + coarse_feats = point_sample(prev_output, points, align_corners=self.align_corners) return coarse_feats - def forward_train(self, inputs, prev_output, img_metas, gt_semantic_seg, - train_cfg): + def forward_train(self, inputs, prev_output, img_metas, gt_semantic_seg, train_cfg): """Forward function for training. Args: inputs (list[Tensor]): List of multi-level img features. @@ -182,18 +171,11 @@ def forward_train(self, inputs, prev_output, img_metas, gt_semantic_seg, """ x = self._transform_inputs(inputs) with torch.no_grad(): - points = self.get_points_train( - prev_output, calculate_uncertainty, cfg=train_cfg) - fine_grained_point_feats = self._get_fine_grained_point_feats( - x, points) + points = self.get_points_train(prev_output, calculate_uncertainty, cfg=train_cfg) + fine_grained_point_feats = self._get_fine_grained_point_feats(x, points) coarse_point_feats = self._get_coarse_point_feats(prev_output, points) - point_logits = self.forward(fine_grained_point_feats, - coarse_point_feats) - point_label = point_sample( - gt_semantic_seg.float(), - points, - mode='nearest', - align_corners=self.align_corners) + point_logits = self.forward(fine_grained_point_feats, coarse_point_feats) + point_label = point_sample(gt_semantic_seg.float(), points, mode="nearest", align_corners=self.align_corners) point_label = point_label.squeeze(1).long() losses = self.losses(point_logits, point_label) @@ -223,34 +205,27 @@ def forward_test(self, inputs, prev_output, img_metas, test_cfg): refined_seg_logits = resize( refined_seg_logits, scale_factor=test_cfg.scale_factor, - mode='bilinear', - align_corners=self.align_corners) + mode="bilinear", + align_corners=self.align_corners, + ) batch_size, channels, height, width = refined_seg_logits.shape - point_indices, points = self.get_points_test( - refined_seg_logits, calculate_uncertainty, cfg=test_cfg) - fine_grained_point_feats = self._get_fine_grained_point_feats( - x, points) - coarse_point_feats = self._get_coarse_point_feats( - prev_output, points) - point_logits = self.forward(fine_grained_point_feats, - coarse_point_feats) + point_indices, points = self.get_points_test(refined_seg_logits, calculate_uncertainty, cfg=test_cfg) + fine_grained_point_feats = self._get_fine_grained_point_feats(x, points) + coarse_point_feats = self._get_coarse_point_feats(prev_output, points) + point_logits = self.forward(fine_grained_point_feats, coarse_point_feats) point_indices = point_indices.unsqueeze(1).expand(-1, channels, -1) - refined_seg_logits = refined_seg_logits.reshape( - batch_size, channels, height * width) - refined_seg_logits = refined_seg_logits.scatter_( - 2, point_indices, point_logits) - refined_seg_logits = refined_seg_logits.view( - batch_size, channels, height, width) + refined_seg_logits = refined_seg_logits.reshape(batch_size, channels, height * width) + refined_seg_logits = refined_seg_logits.scatter_(2, point_indices, point_logits) + refined_seg_logits = refined_seg_logits.view(batch_size, channels, height, width) return refined_seg_logits def losses(self, point_logits, point_label): """Compute segmentation loss.""" loss = dict() - loss['loss_point'] = self.loss_decode( - point_logits, point_label, ignore_index=self.ignore_index) - loss['acc_point'] = accuracy(point_logits, point_label) + loss["loss_point"] = self.loss_decode(point_logits, point_label, ignore_index=self.ignore_index) + loss["acc_point"] = accuracy(point_logits, point_label) return loss def get_points_train(self, seg_logits, uncertainty_func, cfg): @@ -279,8 +254,7 @@ def get_points_train(self, seg_logits, uncertainty_func, cfg): assert 0 <= importance_sample_ratio <= 1 batch_size = seg_logits.shape[0] num_sampled = int(num_points * oversample_ratio) - point_coords = torch.rand( - batch_size, num_sampled, 2, device=seg_logits.device) + point_coords = torch.rand(batch_size, num_sampled, 2, device=seg_logits.device) point_logits = point_sample(seg_logits, point_coords) # It is crucial to calculate uncertainty based on the sampled # prediction value for the points. Calculating uncertainties of the @@ -294,16 +268,12 @@ def get_points_train(self, seg_logits, uncertainty_func, cfg): point_uncertainties = uncertainty_func(point_logits) num_uncertain_points = int(importance_sample_ratio * num_points) num_random_points = num_points - num_uncertain_points - idx = torch.topk( - point_uncertainties[:, 0, :], k=num_uncertain_points, dim=1)[1] - shift = num_sampled * torch.arange( - batch_size, dtype=torch.long, device=seg_logits.device) + idx = torch.topk(point_uncertainties[:, 0, :], k=num_uncertain_points, dim=1)[1] + shift = num_sampled * torch.arange(batch_size, dtype=torch.long, device=seg_logits.device) idx += shift[:, None] - point_coords = point_coords.view(-1, 2)[idx.view(-1), :].view( - batch_size, num_uncertain_points, 2) + point_coords = point_coords.view(-1, 2)[idx.view(-1), :].view(batch_size, num_uncertain_points, 2) if num_random_points > 0: - rand_point_coords = torch.rand( - batch_size, num_random_points, 2, device=seg_logits.device) + rand_point_coords = torch.rand(batch_size, num_random_points, 2, device=seg_logits.device) point_coords = torch.cat((point_coords, rand_point_coords), dim=1) return point_coords @@ -336,14 +306,7 @@ def get_points_test(self, seg_logits, uncertainty_func, cfg): uncertainty_map = uncertainty_map.view(batch_size, height * width) num_points = min(height * width, num_points) point_indices = uncertainty_map.topk(num_points, dim=1)[1] - point_coords = torch.zeros( - batch_size, - num_points, - 2, - dtype=torch.float, - device=seg_logits.device) - point_coords[:, :, 0] = w_step / 2.0 + (point_indices % - width).float() * w_step - point_coords[:, :, 1] = h_step / 2.0 + (point_indices // - width).float() * h_step + point_coords = torch.zeros(batch_size, num_points, 2, dtype=torch.float, device=seg_logits.device) + point_coords[:, :, 0] = w_step / 2.0 + (point_indices % width).float() * w_step + point_coords[:, :, 1] = h_step / 2.0 + (point_indices // width).float() * h_step return point_indices, point_coords diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/psa_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/psa_head.py index 64b694a2..f66a08db 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/psa_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/psa_head.py @@ -1,14 +1,14 @@ import torch import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import HEADS from .decode_head import BaseDecodeHead try: - from comfy_controlnet_preprocessors.uniformer.mmcv.ops import PSAMask + from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.ops import PSAMask except ModuleNotFoundError: PSAMask = None @@ -32,18 +32,20 @@ class PSAHead(BaseDecodeHead): psa_softmax (bool): Whether use softmax for attention. """ - def __init__(self, - mask_size, - psa_type='bi-direction', - compact=False, - shrink_factor=2, - normalization_factor=1.0, - psa_softmax=True, - **kwargs): + def __init__( + self, + mask_size, + psa_type="bi-direction", + compact=False, + shrink_factor=2, + normalization_factor=1.0, + psa_softmax=True, + **kwargs, + ): if PSAMask is None: - raise RuntimeError('Please install mmcv-full for PSAMask ops') + raise RuntimeError("Please install mmcv-full for PSAMask ops") super(PSAHead, self).__init__(**kwargs) - assert psa_type in ['collect', 'distribute', 'bi-direction'] + assert psa_type in ["collect", "distribute", "bi-direction"] self.psa_type = psa_type self.compact = compact self.shrink_factor = shrink_factor @@ -60,7 +62,8 @@ def __init__(self, kernel_size=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.attention = nn.Sequential( ConvModule( self.channels, @@ -68,17 +71,19 @@ def __init__(self, kernel_size=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg), - nn.Conv2d( - self.channels, mask_h * mask_w, kernel_size=1, bias=False)) - if psa_type == 'bi-direction': + act_cfg=self.act_cfg, + ), + nn.Conv2d(self.channels, mask_h * mask_w, kernel_size=1, bias=False), + ) + if psa_type == "bi-direction": self.reduce_p = ConvModule( self.in_channels, self.channels, kernel_size=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.attention_p = nn.Sequential( ConvModule( self.channels, @@ -86,21 +91,23 @@ def __init__(self, kernel_size=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg), - nn.Conv2d( - self.channels, mask_h * mask_w, kernel_size=1, bias=False)) - self.psamask_collect = PSAMask('collect', mask_size) - self.psamask_distribute = PSAMask('distribute', mask_size) + act_cfg=self.act_cfg, + ), + nn.Conv2d(self.channels, mask_h * mask_w, kernel_size=1, bias=False), + ) + self.psamask_collect = PSAMask("collect", mask_size) + self.psamask_distribute = PSAMask("distribute", mask_size) else: self.psamask = PSAMask(psa_type, mask_size) self.proj = ConvModule( - self.channels * (2 if psa_type == 'bi-direction' else 1), + self.channels * (2 if psa_type == "bi-direction" else 1), self.in_channels, kernel_size=1, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) self.bottleneck = ConvModule( self.in_channels * 2, self.channels, @@ -108,14 +115,15 @@ def __init__(self, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) def forward(self, inputs): """Forward function.""" x = self._transform_inputs(inputs) identity = x align_corners = self.align_corners - if self.psa_type in ['collect', 'distribute']: + if self.psa_type in ["collect", "distribute"]: out = self.reduce(x) n, c, h, w = out.size() if self.shrink_factor != 1: @@ -127,23 +135,18 @@ def forward(self, inputs): h = h // self.shrink_factor w = w // self.shrink_factor align_corners = False - out = resize( - out, - size=(h, w), - mode='bilinear', - align_corners=align_corners) + out = resize(out, size=(h, w), mode="bilinear", align_corners=align_corners) y = self.attention(out) if self.compact: - if self.psa_type == 'collect': - y = y.view(n, h * w, - h * w).transpose(1, 2).view(n, h * w, h, w) + if self.psa_type == "collect": + y = y.view(n, h * w, h * w).transpose(1, 2).view(n, h * w, h, w) else: y = self.psamask(y) if self.psa_softmax: y = F.softmax(y, dim=1) - out = torch.bmm( - out.view(n, c, h * w), y.view(n, h * w, h * w)).view( - n, c, h, w) * (1.0 / self.normalization_factor) + out = torch.bmm(out.view(n, c, h * w), y.view(n, h * w, h * w)).view(n, c, h, w) * ( + 1.0 / self.normalization_factor + ) else: x_col = self.reduce(x) x_dis = self.reduce_p(x) @@ -157,40 +160,27 @@ def forward(self, inputs): h = h // self.shrink_factor w = w // self.shrink_factor align_corners = False - x_col = resize( - x_col, - size=(h, w), - mode='bilinear', - align_corners=align_corners) - x_dis = resize( - x_dis, - size=(h, w), - mode='bilinear', - align_corners=align_corners) + x_col = resize(x_col, size=(h, w), mode="bilinear", align_corners=align_corners) + x_dis = resize(x_dis, size=(h, w), mode="bilinear", align_corners=align_corners) y_col = self.attention(x_col) y_dis = self.attention_p(x_dis) if self.compact: - y_dis = y_dis.view(n, h * w, - h * w).transpose(1, 2).view(n, h * w, h, w) + y_dis = y_dis.view(n, h * w, h * w).transpose(1, 2).view(n, h * w, h, w) else: y_col = self.psamask_collect(y_col) y_dis = self.psamask_distribute(y_dis) if self.psa_softmax: y_col = F.softmax(y_col, dim=1) y_dis = F.softmax(y_dis, dim=1) - x_col = torch.bmm( - x_col.view(n, c, h * w), y_col.view(n, h * w, h * w)).view( - n, c, h, w) * (1.0 / self.normalization_factor) - x_dis = torch.bmm( - x_dis.view(n, c, h * w), y_dis.view(n, h * w, h * w)).view( - n, c, h, w) * (1.0 / self.normalization_factor) + x_col = torch.bmm(x_col.view(n, c, h * w), y_col.view(n, h * w, h * w)).view(n, c, h, w) * ( + 1.0 / self.normalization_factor + ) + x_dis = torch.bmm(x_dis.view(n, c, h * w), y_dis.view(n, h * w, h * w)).view(n, c, h, w) * ( + 1.0 / self.normalization_factor + ) out = torch.cat([x_col, x_dis], 1) out = self.proj(out) - out = resize( - out, - size=identity.shape[2:], - mode='bilinear', - align_corners=align_corners) + out = resize(out, size=identity.shape[2:], mode="bilinear", align_corners=align_corners) out = self.bottleneck(torch.cat((identity, out), dim=1)) out = self.cls_seg(out) return out diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/psp_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/psp_head.py index e741e1c9..a25910cd 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/psp_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/psp_head.py @@ -1,8 +1,8 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import HEADS from .decode_head import BaseDecodeHead @@ -21,8 +21,7 @@ class PPM(nn.ModuleList): align_corners (bool): align_corners argument of F.interpolate. """ - def __init__(self, pool_scales, in_channels, channels, conv_cfg, norm_cfg, - act_cfg, align_corners): + def __init__(self, pool_scales, in_channels, channels, conv_cfg, norm_cfg, act_cfg, align_corners): super(PPM, self).__init__() self.pool_scales = pool_scales self.align_corners = align_corners @@ -41,18 +40,17 @@ def __init__(self, pool_scales, in_channels, channels, conv_cfg, norm_cfg, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg))) + act_cfg=self.act_cfg, + ), + ) + ) def forward(self, x): """Forward function.""" ppm_outs = [] for ppm in self: ppm_out = ppm(x) - upsampled_ppm_out = resize( - ppm_out, - size=x.size()[2:], - mode='bilinear', - align_corners=self.align_corners) + upsampled_ppm_out = resize(ppm_out, size=x.size()[2:], mode="bilinear", align_corners=self.align_corners) ppm_outs.append(upsampled_ppm_out) return ppm_outs @@ -80,7 +78,8 @@ def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs): conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg, - align_corners=self.align_corners) + align_corners=self.align_corners, + ) self.bottleneck = ConvModule( self.in_channels + len(pool_scales) * self.channels, self.channels, @@ -88,7 +87,8 @@ def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) def forward(self, inputs): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/sep_aspp_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/sep_aspp_head.py index b1090f5d..b53de790 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/sep_aspp_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/sep_aspp_head.py @@ -1,8 +1,8 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, DepthwiseSeparableConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, DepthwiseSeparableConvModule -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import HEADS from .aspp_head import ASPPHead, ASPPModule @@ -22,7 +22,8 @@ def __init__(self, **kwargs): dilation=dilation, padding=dilation, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) @HEADS.register_module() @@ -48,53 +49,33 @@ def __init__(self, c1_in_channels, c1_channels, **kwargs): channels=self.channels, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) if c1_in_channels > 0: self.c1_bottleneck = ConvModule( - c1_in_channels, - c1_channels, - 1, - conv_cfg=self.conv_cfg, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + c1_in_channels, c1_channels, 1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ) else: self.c1_bottleneck = None self.sep_bottleneck = nn.Sequential( DepthwiseSeparableConvModule( - self.channels + c1_channels, - self.channels, - 3, - padding=1, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg), + self.channels + c1_channels, self.channels, 3, padding=1, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ), DepthwiseSeparableConvModule( - self.channels, - self.channels, - 3, - padding=1, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg)) + self.channels, self.channels, 3, padding=1, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg + ), + ) def forward(self, inputs): """Forward function.""" x = self._transform_inputs(inputs) - aspp_outs = [ - resize( - self.image_pool(x), - size=x.size()[2:], - mode='bilinear', - align_corners=self.align_corners) - ] + aspp_outs = [resize(self.image_pool(x), size=x.size()[2:], mode="bilinear", align_corners=self.align_corners)] aspp_outs.extend(self.aspp_modules(x)) aspp_outs = torch.cat(aspp_outs, dim=1) output = self.bottleneck(aspp_outs) if self.c1_bottleneck is not None: c1_output = self.c1_bottleneck(inputs[0]) - output = resize( - input=output, - size=c1_output.shape[2:], - mode='bilinear', - align_corners=self.align_corners) + output = resize(input=output, size=c1_output.shape[2:], mode="bilinear", align_corners=self.align_corners) output = torch.cat([output, c1_output], dim=1) output = self.sep_bottleneck(output) output = self.cls_seg(output) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/sep_fcn_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/sep_fcn_head.py index d4a9c776..93ef50a1 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/sep_fcn_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/sep_fcn_head.py @@ -1,4 +1,4 @@ -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import DepthwiseSeparableConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import DepthwiseSeparableConvModule from ..builder import HEADS from .fcn_head import FCNHead @@ -33,14 +33,16 @@ def __init__(self, **kwargs): self.channels, kernel_size=self.kernel_size, padding=self.kernel_size // 2, - norm_cfg=self.norm_cfg) + norm_cfg=self.norm_cfg, + ) for i in range(1, self.num_convs): self.convs[i] = DepthwiseSeparableConvModule( self.channels, self.channels, kernel_size=self.kernel_size, padding=self.kernel_size // 2, - norm_cfg=self.norm_cfg) + norm_cfg=self.norm_cfg, + ) if self.concat_input: self.conv_cat = DepthwiseSeparableConvModule( @@ -48,4 +50,5 @@ def __init__(self, **kwargs): self.channels, kernel_size=self.kernel_size, padding=self.kernel_size // 2, - norm_cfg=self.norm_cfg) + norm_cfg=self.norm_cfg, + ) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/uper_head.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/uper_head.py index 46d4ff97..7d1ffbcc 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/uper_head.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/decode_heads/uper_head.py @@ -1,8 +1,8 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from ..builder import HEADS from .decode_head import BaseDecodeHead from .psp_head import PPM @@ -21,8 +21,7 @@ class UPerHead(BaseDecodeHead): """ def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs): - super(UPerHead, self).__init__( - input_transform='multiple_select', **kwargs) + super(UPerHead, self).__init__(input_transform="multiple_select", **kwargs) # PSP Module self.psp_modules = PPM( pool_scales, @@ -31,7 +30,8 @@ def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs): conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg, - align_corners=self.align_corners) + align_corners=self.align_corners, + ) self.bottleneck = ConvModule( self.in_channels[-1] + len(pool_scales) * self.channels, self.channels, @@ -39,7 +39,8 @@ def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) # FPN Module self.lateral_convs = nn.ModuleList() self.fpn_convs = nn.ModuleList() @@ -51,7 +52,8 @@ def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs): conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg, - inplace=False) + inplace=False, + ) fpn_conv = ConvModule( self.channels, self.channels, @@ -60,7 +62,8 @@ def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs): conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg, - inplace=False) + inplace=False, + ) self.lateral_convs.append(l_conv) self.fpn_convs.append(fpn_conv) @@ -71,7 +74,8 @@ def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs): padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg) + act_cfg=self.act_cfg, + ) def psp_forward(self, inputs): """Forward function of PSP module.""" @@ -89,10 +93,7 @@ def forward(self, inputs): inputs = self._transform_inputs(inputs) # build laterals - laterals = [ - lateral_conv(inputs[i]) - for i, lateral_conv in enumerate(self.lateral_convs) - ] + laterals = [lateral_conv(inputs[i]) for i, lateral_conv in enumerate(self.lateral_convs)] laterals.append(self.psp_forward(inputs)) @@ -100,26 +101,17 @@ def forward(self, inputs): used_backbone_levels = len(laterals) for i in range(used_backbone_levels - 1, 0, -1): prev_shape = laterals[i - 1].shape[2:] - laterals[i - 1] += resize( - laterals[i], - size=prev_shape, - mode='bilinear', - align_corners=self.align_corners) + laterals[i - 1] += resize(laterals[i], size=prev_shape, mode="bilinear", align_corners=self.align_corners) # build outputs - fpn_outs = [ - self.fpn_convs[i](laterals[i]) - for i in range(used_backbone_levels - 1) - ] + fpn_outs = [self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels - 1)] # append psp feature fpn_outs.append(laterals[-1]) for i in range(used_backbone_levels - 1, 0, -1): fpn_outs[i] = resize( - fpn_outs[i], - size=fpn_outs[0].shape[2:], - mode='bilinear', - align_corners=self.align_corners) + fpn_outs[i], size=fpn_outs[0].shape[2:], mode="bilinear", align_corners=self.align_corners + ) fpn_outs = torch.cat(fpn_outs, dim=1) output = self.fpn_bottleneck(fpn_outs) output = self.cls_seg(output) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/losses/lovasz_loss.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/losses/lovasz_loss.py index 43c9c2b4..b72b8887 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/losses/lovasz_loss.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/losses/lovasz_loss.py @@ -2,7 +2,7 @@ ch/lovasz_losses.py Lovasz-Softmax and Jaccard hinge loss in PyTorch Maxim Berman 2018 ESAT-PSI KU Leuven (MIT License)""" -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import torch import torch.nn as nn import torch.nn.functional as F @@ -20,7 +20,7 @@ def lovasz_grad(gt_sorted): gts = gt_sorted.sum() intersection = gts - gt_sorted.float().cumsum(0) union = gts + (1 - gt_sorted).float().cumsum(0) - jaccard = 1. - intersection / union + jaccard = 1.0 - intersection / union if p > 1: # cover 1-pixel case jaccard[1:p] = jaccard[1:p] - jaccard[0:-1] return jaccard @@ -33,7 +33,7 @@ def flatten_binary_logits(logits, labels, ignore_index=None): labels = labels.view(-1) if ignore_index is None: return logits, labels - valid = (labels != ignore_index) + valid = labels != ignore_index vlogits = logits[valid] vlabels = labels[valid] return vlogits, vlabels @@ -50,7 +50,7 @@ def flatten_probs(probs, labels, ignore_index=None): labels = labels.view(-1) if ignore_index is None: return probs, labels - valid = (labels != ignore_index) + valid = labels != ignore_index vprobs = probs[valid.nonzero().squeeze()] vlabels = labels[valid] return vprobs, vlabels @@ -69,9 +69,9 @@ def lovasz_hinge_flat(logits, labels): """ if len(labels) == 0: # only void pixels, the gradients should be 0 - return logits.sum() * 0. - signs = 2. * labels.float() - 1. - errors = (1. - logits * signs) + return logits.sum() * 0.0 + signs = 2.0 * labels.float() - 1.0 + errors = 1.0 - logits * signs errors_sorted, perm = torch.sort(errors, dim=0, descending=True) perm = perm.data gt_sorted = labels[perm] @@ -80,14 +80,16 @@ def lovasz_hinge_flat(logits, labels): return loss -def lovasz_hinge(logits, - labels, - classes='present', - per_image=False, - class_weight=None, - reduction='mean', - avg_factor=None, - ignore_index=255): +def lovasz_hinge( + logits, + labels, + classes="present", + per_image=False, + class_weight=None, + reduction="mean", + avg_factor=None, + ignore_index=255, +): """Binary Lovasz hinge loss. Args: @@ -113,19 +115,16 @@ def lovasz_hinge(logits, """ if per_image: loss = [ - lovasz_hinge_flat(*flatten_binary_logits( - logit.unsqueeze(0), label.unsqueeze(0), ignore_index)) + lovasz_hinge_flat(*flatten_binary_logits(logit.unsqueeze(0), label.unsqueeze(0), ignore_index)) for logit, label in zip(logits, labels) ] - loss = weight_reduce_loss( - torch.stack(loss), None, reduction, avg_factor) + loss = weight_reduce_loss(torch.stack(loss), None, reduction, avg_factor) else: - loss = lovasz_hinge_flat( - *flatten_binary_logits(logits, labels, ignore_index)) + loss = lovasz_hinge_flat(*flatten_binary_logits(logits, labels, ignore_index)) return loss -def lovasz_softmax_flat(probs, labels, classes='present', class_weight=None): +def lovasz_softmax_flat(probs, labels, classes="present", class_weight=None): """Multi-class Lovasz-Softmax loss. Args: @@ -143,17 +142,17 @@ def lovasz_softmax_flat(probs, labels, classes='present', class_weight=None): """ if probs.numel() == 0: # only void pixels, the gradients should be 0 - return probs * 0. + return probs * 0.0 C = probs.size(1) losses = [] - class_to_sum = list(range(C)) if classes in ['all', 'present'] else classes + class_to_sum = list(range(C)) if classes in ["all", "present"] else classes for c in class_to_sum: fg = (labels == c).float() # foreground for class c - if (classes == 'present' and fg.sum() == 0): + if classes == "present" and fg.sum() == 0: continue if C == 1: if len(classes) > 1: - raise ValueError('Sigmoid output possible only with 1 class') + raise ValueError("Sigmoid output possible only with 1 class") class_pred = probs[:, 0] else: class_pred = probs[:, c] @@ -168,14 +167,16 @@ def lovasz_softmax_flat(probs, labels, classes='present', class_weight=None): return torch.stack(losses).mean() -def lovasz_softmax(probs, - labels, - classes='present', - per_image=False, - class_weight=None, - reduction='mean', - avg_factor=None, - ignore_index=255): +def lovasz_softmax( + probs, + labels, + classes="present", + per_image=False, + class_weight=None, + reduction="mean", + avg_factor=None, + ignore_index=255, +): """Multi-class Lovasz-Softmax loss. Args: @@ -205,19 +206,17 @@ def lovasz_softmax(probs, if per_image: loss = [ lovasz_softmax_flat( - *flatten_probs( - prob.unsqueeze(0), label.unsqueeze(0), ignore_index), + *flatten_probs(prob.unsqueeze(0), label.unsqueeze(0), ignore_index), classes=classes, - class_weight=class_weight) + class_weight=class_weight, + ) for prob, label in zip(probs, labels) ] - loss = weight_reduce_loss( - torch.stack(loss), None, reduction, avg_factor) + loss = weight_reduce_loss(torch.stack(loss), None, reduction, avg_factor) else: loss = lovasz_softmax_flat( - *flatten_probs(probs, labels, ignore_index), - classes=classes, - class_weight=class_weight) + *flatten_probs(probs, labels, ignore_index), classes=classes, class_weight=class_weight + ) return loss @@ -245,24 +244,31 @@ class LovaszLoss(nn.Module): loss_weight (float, optional): Weight of the loss. Defaults to 1.0. """ - def __init__(self, - loss_type='multi_class', - classes='present', - per_image=False, - reduction='mean', - class_weight=None, - loss_weight=1.0): + def __init__( + self, + loss_type="multi_class", + classes="present", + per_image=False, + reduction="mean", + class_weight=None, + loss_weight=1.0, + ): super(LovaszLoss, self).__init__() - assert loss_type in ('binary', 'multi_class'), "loss_type should be \ + assert loss_type in ( + "binary", + "multi_class", + ), "loss_type should be \ 'binary' or 'multi_class'." - if loss_type == 'binary': + if loss_type == "binary": self.cls_criterion = lovasz_hinge else: self.cls_criterion = lovasz_softmax - assert classes in ('all', 'present') or mmcv.is_list_of(classes, int) + assert classes in ("all", "present") or mmcv.is_list_of(classes, int) if not per_image: - assert reduction == 'none', "reduction should be 'none' when \ + assert ( + reduction == "none" + ), "reduction should be 'none' when \ per_image is False." self.classes = classes @@ -271,17 +277,10 @@ def __init__(self, self.loss_weight = loss_weight self.class_weight = get_class_weight(class_weight) - def forward(self, - cls_score, - label, - weight=None, - avg_factor=None, - reduction_override=None, - **kwargs): + def forward(self, cls_score, label, weight=None, avg_factor=None, reduction_override=None, **kwargs): """Forward function.""" - assert reduction_override in (None, 'none', 'mean', 'sum') - reduction = ( - reduction_override if reduction_override else self.reduction) + assert reduction_override in (None, "none", "mean", "sum") + reduction = reduction_override if reduction_override else self.reduction if self.class_weight is not None: class_weight = cls_score.new_tensor(self.class_weight) else: @@ -299,5 +298,6 @@ def forward(self, class_weight=class_weight, reduction=reduction, avg_factor=avg_factor, - **kwargs) + **kwargs, + ) return loss_cls diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/losses/utils.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/losses/utils.py index b19cd149..564aab6d 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/losses/utils.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/losses/utils.py @@ -1,6 +1,6 @@ import functools -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import numpy as np import torch.nn.functional as F @@ -14,7 +14,7 @@ def get_class_weight(class_weight): """ if isinstance(class_weight, str): # take it as a file path - if class_weight.endswith('.npy'): + if class_weight.endswith(".npy"): class_weight = np.load(class_weight) else: # pkl, json or yaml @@ -43,7 +43,7 @@ def reduce_loss(loss, reduction): return loss.sum() -def weight_reduce_loss(loss, weight=None, reduction='mean', avg_factor=None): +def weight_reduce_loss(loss, weight=None, reduction="mean", avg_factor=None): """Apply element-wise weight and reduce loss. Args: @@ -67,10 +67,10 @@ def weight_reduce_loss(loss, weight=None, reduction='mean', avg_factor=None): loss = reduce_loss(loss, reduction) else: # if reduction is mean, then average the loss by avg_factor - if reduction == 'mean': + if reduction == "mean": loss = loss.sum() / avg_factor # if reduction is 'none', then do nothing, otherwise raise an error - elif reduction != 'none': + elif reduction != "none": raise ValueError('avg_factor can not be used with reduction="sum"') return loss @@ -107,12 +107,7 @@ def weighted_loss(loss_func): """ @functools.wraps(loss_func) - def wrapper(pred, - target, - weight=None, - reduction='mean', - avg_factor=None, - **kwargs): + def wrapper(pred, target, weight=None, reduction="mean", avg_factor=None, **kwargs): # get element-wise loss loss = loss_func(pred, target, **kwargs) loss = weight_reduce_loss(loss, weight, reduction, avg_factor) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/necks/fpn.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/necks/fpn.py index 533445ee..6143b2d1 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/necks/fpn.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/necks/fpn.py @@ -1,6 +1,6 @@ import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, xavier_init +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, xavier_init from ..builder import NECKS @@ -60,20 +60,22 @@ class FPN(nn.Module): outputs[3].shape = torch.Size([1, 11, 43, 43]) """ - def __init__(self, - in_channels, - out_channels, - num_outs, - start_level=0, - end_level=-1, - add_extra_convs=False, - extra_convs_on_inputs=False, - relu_before_extra_convs=False, - no_norm_on_lateral=False, - conv_cfg=None, - norm_cfg=None, - act_cfg=None, - upsample_cfg=dict(mode='nearest')): + def __init__( + self, + in_channels, + out_channels, + num_outs, + start_level=0, + end_level=-1, + add_extra_convs=False, + extra_convs_on_inputs=False, + relu_before_extra_convs=False, + no_norm_on_lateral=False, + conv_cfg=None, + norm_cfg=None, + act_cfg=None, + upsample_cfg=dict(mode="nearest"), + ): super(FPN, self).__init__() assert isinstance(in_channels, list) self.in_channels = in_channels @@ -99,14 +101,14 @@ def __init__(self, assert isinstance(add_extra_convs, (str, bool)) if isinstance(add_extra_convs, str): # Extra_convs_source choices: 'on_input', 'on_lateral', 'on_output' - assert add_extra_convs in ('on_input', 'on_lateral', 'on_output') + assert add_extra_convs in ("on_input", "on_lateral", "on_output") elif add_extra_convs: # True if extra_convs_on_inputs: # For compatibility with previous release # TODO: deprecate `extra_convs_on_inputs` - self.add_extra_convs = 'on_input' + self.add_extra_convs = "on_input" else: - self.add_extra_convs = 'on_output' + self.add_extra_convs = "on_output" self.lateral_convs = nn.ModuleList() self.fpn_convs = nn.ModuleList() @@ -119,7 +121,8 @@ def __init__(self, conv_cfg=conv_cfg, norm_cfg=norm_cfg if not self.no_norm_on_lateral else None, act_cfg=act_cfg, - inplace=False) + inplace=False, + ) fpn_conv = ConvModule( out_channels, out_channels, @@ -128,7 +131,8 @@ def __init__(self, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg, - inplace=False) + inplace=False, + ) self.lateral_convs.append(l_conv) self.fpn_convs.append(fpn_conv) @@ -137,7 +141,7 @@ def __init__(self, extra_levels = num_outs - self.backbone_end_level + self.start_level if self.add_extra_convs and extra_levels >= 1: for i in range(extra_levels): - if i == 0 and self.add_extra_convs == 'on_input': + if i == 0 and self.add_extra_convs == "on_input": in_channels = self.in_channels[self.backbone_end_level - 1] else: in_channels = out_channels @@ -150,42 +154,36 @@ def __init__(self, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg, - inplace=False) + inplace=False, + ) self.fpn_convs.append(extra_fpn_conv) # default init_weights for conv(msra) and norm in ConvModule def init_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): - xavier_init(m, distribution='uniform') + xavier_init(m, distribution="uniform") def forward(self, inputs): assert len(inputs) == len(self.in_channels) # build laterals - laterals = [ - lateral_conv(inputs[i + self.start_level]) - for i, lateral_conv in enumerate(self.lateral_convs) - ] + laterals = [lateral_conv(inputs[i + self.start_level]) for i, lateral_conv in enumerate(self.lateral_convs)] # build top-down path used_backbone_levels = len(laterals) for i in range(used_backbone_levels - 1, 0, -1): # In some cases, fixing `scale factor` (e.g. 2) is preferred, but # it cannot co-exist with `size` in `F.interpolate`. - if 'scale_factor' in self.upsample_cfg: - laterals[i - 1] += F.interpolate(laterals[i], - **self.upsample_cfg) + if "scale_factor" in self.upsample_cfg: + laterals[i - 1] += F.interpolate(laterals[i], **self.upsample_cfg) else: prev_shape = laterals[i - 1].shape[2:] - laterals[i - 1] += F.interpolate( - laterals[i], size=prev_shape, **self.upsample_cfg) + laterals[i - 1] += F.interpolate(laterals[i], size=prev_shape, **self.upsample_cfg) # build outputs # part 1: from original levels - outs = [ - self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels) - ] + outs = [self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)] # part 2: add extra levels if self.num_outs > len(outs): # use max pool to get more levels on top of outputs @@ -195,11 +193,11 @@ def forward(self, inputs): outs.append(F.max_pool2d(outs[-1], 1, stride=2)) # add conv layers on top of original feature maps (RetinaNet) else: - if self.add_extra_convs == 'on_input': + if self.add_extra_convs == "on_input": extra_source = inputs[self.backbone_end_level - 1] - elif self.add_extra_convs == 'on_lateral': + elif self.add_extra_convs == "on_lateral": extra_source = laterals[-1] - elif self.add_extra_convs == 'on_output': + elif self.add_extra_convs == "on_output": extra_source = outs[-1] else: raise NotImplementedError diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/necks/multilevel_neck.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/necks/multilevel_neck.py index 30460932..eac07317 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/necks/multilevel_neck.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/necks/multilevel_neck.py @@ -1,6 +1,6 @@ import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule from ..builder import NECKS @@ -19,12 +19,7 @@ class MultiLevelNeck(nn.Module): Default: None. """ - def __init__(self, - in_channels, - out_channels, - scales=[0.5, 1, 2, 4], - norm_cfg=None, - act_cfg=None): + def __init__(self, in_channels, out_channels, scales=[0.5, 1, 2, 4], norm_cfg=None, act_cfg=None): super(MultiLevelNeck, self).__init__() assert isinstance(in_channels, list) self.in_channels = in_channels @@ -35,36 +30,24 @@ def __init__(self, self.convs = nn.ModuleList() for in_channel in in_channels: self.lateral_convs.append( - ConvModule( - in_channel, - out_channels, - kernel_size=1, - norm_cfg=norm_cfg, - act_cfg=act_cfg)) + ConvModule(in_channel, out_channels, kernel_size=1, norm_cfg=norm_cfg, act_cfg=act_cfg) + ) for _ in range(self.num_outs): self.convs.append( ConvModule( - out_channels, - out_channels, - kernel_size=3, - padding=1, - stride=1, - norm_cfg=norm_cfg, - act_cfg=act_cfg)) + out_channels, out_channels, kernel_size=3, padding=1, stride=1, norm_cfg=norm_cfg, act_cfg=act_cfg + ) + ) def forward(self, inputs): assert len(inputs) == len(self.in_channels) print(inputs[0].shape) - inputs = [ - lateral_conv(inputs[i]) - for i, lateral_conv in enumerate(self.lateral_convs) - ] + inputs = [lateral_conv(inputs[i]) for i, lateral_conv in enumerate(self.lateral_convs)] # for len(inputs) not equal to self.num_outs if len(inputs) == 1: inputs = [inputs[0] for _ in range(self.num_outs)] outs = [] for i in range(self.num_outs): - x_resize = F.interpolate( - inputs[i], scale_factor=self.scales[i], mode='bilinear') + x_resize = F.interpolate(inputs[i], scale_factor=self.scales[i], mode="bilinear") outs.append(self.convs[i](x_resize)) return tuple(outs) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/base.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/base.py index 4dcaba6b..f9a9009f 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/base.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/base.py @@ -3,12 +3,12 @@ from abc import ABCMeta, abstractmethod from collections import OrderedDict -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import numpy as np import torch import torch.distributed as dist import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.runner import auto_fp16 +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.runner import auto_fp16 class BaseSegmentor(nn.Module): @@ -23,18 +23,17 @@ def __init__(self): @property def with_neck(self): """bool: whether the segmentor has neck""" - return hasattr(self, 'neck') and self.neck is not None + return hasattr(self, "neck") and self.neck is not None @property def with_auxiliary_head(self): """bool: whether the segmentor has auxiliary head""" - return hasattr(self, - 'auxiliary_head') and self.auxiliary_head is not None + return hasattr(self, "auxiliary_head") and self.auxiliary_head is not None @property def with_decode_head(self): """bool: whether the segmentor has decode head""" - return hasattr(self, 'decode_head') and self.decode_head is not None + return hasattr(self, "decode_head") and self.decode_head is not None @abstractmethod def extract_feat(self, imgs): @@ -71,7 +70,7 @@ def init_weights(self, pretrained=None): """ if pretrained is not None: logger = logging.getLogger() - logger.info(f'load model from: {pretrained}') + logger.info(f"load model from: {pretrained}") def forward_test(self, imgs, img_metas, **kwargs): """ @@ -83,23 +82,21 @@ def forward_test(self, imgs, img_metas, **kwargs): augs (multiscale, flip, etc.) and the inner list indicates images in a batch. """ - for var, name in [(imgs, 'imgs'), (img_metas, 'img_metas')]: + for var, name in [(imgs, "imgs"), (img_metas, "img_metas")]: if not isinstance(var, list): - raise TypeError(f'{name} must be a list, but got ' - f'{type(var)}') + raise TypeError(f"{name} must be a list, but got " f"{type(var)}") num_augs = len(imgs) if num_augs != len(img_metas): - raise ValueError(f'num of augmentations ({len(imgs)}) != ' - f'num of image meta ({len(img_metas)})') + raise ValueError(f"num of augmentations ({len(imgs)}) != " f"num of image meta ({len(img_metas)})") # all images in the same aug batch all of the same ori_shape and pad # shape for img_meta in img_metas: - ori_shapes = [_['ori_shape'] for _ in img_meta] + ori_shapes = [_["ori_shape"] for _ in img_meta] assert all(shape == ori_shapes[0] for shape in ori_shapes) - img_shapes = [_['img_shape'] for _ in img_meta] + img_shapes = [_["img_shape"] for _ in img_meta] assert all(shape == img_shapes[0] for shape in img_shapes) - pad_shapes = [_['pad_shape'] for _ in img_meta] + pad_shapes = [_["pad_shape"] for _ in img_meta] assert all(shape == pad_shapes[0] for shape in pad_shapes) if num_augs == 1: @@ -107,7 +104,7 @@ def forward_test(self, imgs, img_metas, **kwargs): else: return self.aug_test(imgs, img_metas, **kwargs) - @auto_fp16(apply_to=('img', )) + @auto_fp16(apply_to=("img",)) def forward(self, img, img_metas, return_loss=True, **kwargs): """Calls either :func:`forward_train` or :func:`forward_test` depending on whether ``return_loss`` is ``True``. @@ -152,10 +149,7 @@ def train_step(self, data_batch, optimizer, **kwargs): losses = self(**data_batch) loss, log_vars = self._parse_losses(losses) - outputs = dict( - loss=loss, - log_vars=log_vars, - num_samples=len(data_batch['img_metas'])) + outputs = dict(loss=loss, log_vars=log_vars, num_samples=len(data_batch["img_metas"])) return outputs @@ -189,13 +183,11 @@ def _parse_losses(losses): elif isinstance(loss_value, list): log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value) else: - raise TypeError( - f'{loss_name} is not a tensor or list of tensors') + raise TypeError(f"{loss_name} is not a tensor or list of tensors") - loss = sum(_value for _key, _value in log_vars.items() - if 'loss' in _key) + loss = sum(_value for _key, _value in log_vars.items() if "loss" in _key) - log_vars['loss'] = loss + log_vars["loss"] = loss for loss_name, loss_value in log_vars.items(): # reduce loss when distributed training if dist.is_available() and dist.is_initialized(): @@ -205,15 +197,7 @@ def _parse_losses(losses): return loss, log_vars - def show_result(self, - img, - result, - palette=None, - win_name='', - show=False, - wait_time=0, - out_file=None, - opacity=0.5): + def show_result(self, img, result, palette=None, win_name="", show=False, wait_time=0, out_file=None, opacity=0.5): """Draw `result` over `img`. Args: @@ -241,8 +225,7 @@ def show_result(self, seg = result[0] if palette is None: if self.PALETTE is None: - palette = np.random.randint( - 0, 255, size=(len(self.CLASSES), 3)) + palette = np.random.randint(0, 255, size=(len(self.CLASSES), 3)) else: palette = self.PALETTE palette = np.array(palette) @@ -268,6 +251,5 @@ def show_result(self, mmcv.imwrite(img, out_file) if not (show or out_file): - warnings.warn('show==False and out_file is not specified, only ' - 'result image will be returned') + warnings.warn("show==False and out_file is not specified, only " "result image will be returned") return img diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/cascade_encoder_decoder.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/cascade_encoder_decoder.py index 3a9bdb25..93d35420 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/cascade_encoder_decoder.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/cascade_encoder_decoder.py @@ -1,7 +1,7 @@ from torch import nn -from comfy_controlnet_preprocessors.uniformer.mmseg.core import add_prefix -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.core import add_prefix +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from .. import builder from ..builder import SEGMENTORS from .encoder_decoder import EncoderDecoder @@ -16,15 +16,17 @@ class CascadeEncoderDecoder(EncoderDecoder): will be the input of next decoder_head. """ - def __init__(self, - num_stages, - backbone, - decode_head, - neck=None, - auxiliary_head=None, - train_cfg=None, - test_cfg=None, - pretrained=None): + def __init__( + self, + num_stages, + backbone, + decode_head, + neck=None, + auxiliary_head=None, + train_cfg=None, + test_cfg=None, + pretrained=None, + ): self.num_stages = num_stages super(CascadeEncoderDecoder, self).__init__( backbone=backbone, @@ -33,7 +35,8 @@ def __init__(self, auxiliary_head=auxiliary_head, train_cfg=train_cfg, test_cfg=test_cfg, - pretrained=pretrained) + pretrained=pretrained, + ) def _init_decode_head(self, decode_head): """Initialize ``decode_head``""" @@ -68,13 +71,8 @@ def encode_decode(self, img, img_metas): x = self.extract_feat(img) out = self.decode_head[0].forward_test(x, img_metas, self.test_cfg) for i in range(1, self.num_stages): - out = self.decode_head[i].forward_test(x, out, img_metas, - self.test_cfg) - out = resize( - input=out, - size=img.shape[2:], - mode='bilinear', - align_corners=self.align_corners) + out = self.decode_head[i].forward_test(x, out, img_metas, self.test_cfg) + out = resize(input=out, size=img.shape[2:], mode="bilinear", align_corners=self.align_corners) return out def _decode_head_forward_train(self, x, img_metas, gt_semantic_seg): @@ -82,17 +80,16 @@ def _decode_head_forward_train(self, x, img_metas, gt_semantic_seg): training.""" losses = dict() - loss_decode = self.decode_head[0].forward_train( - x, img_metas, gt_semantic_seg, self.train_cfg) + loss_decode = self.decode_head[0].forward_train(x, img_metas, gt_semantic_seg, self.train_cfg) - losses.update(add_prefix(loss_decode, 'decode_0')) + losses.update(add_prefix(loss_decode, "decode_0")) for i in range(1, self.num_stages): # forward test again, maybe unnecessary for most methods. - prev_outputs = self.decode_head[i - 1].forward_test( - x, img_metas, self.test_cfg) + prev_outputs = self.decode_head[i - 1].forward_test(x, img_metas, self.test_cfg) loss_decode = self.decode_head[i].forward_train( - x, prev_outputs, img_metas, gt_semantic_seg, self.train_cfg) - losses.update(add_prefix(loss_decode, f'decode_{i}')) + x, prev_outputs, img_metas, gt_semantic_seg, self.train_cfg + ) + losses.update(add_prefix(loss_decode, f"decode_{i}")) return losses diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/encoder_decoder.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/encoder_decoder.py index 2f02dcca..8deb6497 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/encoder_decoder.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/segmentors/encoder_decoder.py @@ -2,8 +2,8 @@ import torch.nn as nn import torch.nn.functional as F -from comfy_controlnet_preprocessors.uniformer.mmseg.core import add_prefix -from comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.core import add_prefix +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg.ops import resize from .. import builder from ..builder import SEGMENTORS from .base import BaseSegmentor @@ -18,14 +18,9 @@ class EncoderDecoder(BaseSegmentor): which could be dumped during inference. """ - def __init__(self, - backbone, - decode_head, - neck=None, - auxiliary_head=None, - train_cfg=None, - test_cfg=None, - pretrained=None): + def __init__( + self, backbone, decode_head, neck=None, auxiliary_head=None, train_cfg=None, test_cfg=None, pretrained=None + ): super(EncoderDecoder, self).__init__() self.backbone = builder.build_backbone(backbone) if neck is not None: @@ -86,22 +81,16 @@ def encode_decode(self, img, img_metas): map of the same size as input.""" x = self.extract_feat(img) out = self._decode_head_forward_test(x, img_metas) - out = resize( - input=out, - size=img.shape[2:], - mode='bilinear', - align_corners=self.align_corners) + out = resize(input=out, size=img.shape[2:], mode="bilinear", align_corners=self.align_corners) return out def _decode_head_forward_train(self, x, img_metas, gt_semantic_seg): """Run forward function and calculate loss for decode head in training.""" losses = dict() - loss_decode = self.decode_head.forward_train(x, img_metas, - gt_semantic_seg, - self.train_cfg) + loss_decode = self.decode_head.forward_train(x, img_metas, gt_semantic_seg, self.train_cfg) - losses.update(add_prefix(loss_decode, 'decode')) + losses.update(add_prefix(loss_decode, "decode")) return losses def _decode_head_forward_test(self, x, img_metas): @@ -116,14 +105,11 @@ def _auxiliary_head_forward_train(self, x, img_metas, gt_semantic_seg): losses = dict() if isinstance(self.auxiliary_head, nn.ModuleList): for idx, aux_head in enumerate(self.auxiliary_head): - loss_aux = aux_head.forward_train(x, img_metas, - gt_semantic_seg, - self.train_cfg) - losses.update(add_prefix(loss_aux, f'aux_{idx}')) + loss_aux = aux_head.forward_train(x, img_metas, gt_semantic_seg, self.train_cfg) + losses.update(add_prefix(loss_aux, f"aux_{idx}")) else: - loss_aux = self.auxiliary_head.forward_train( - x, img_metas, gt_semantic_seg, self.train_cfg) - losses.update(add_prefix(loss_aux, 'aux')) + loss_aux = self.auxiliary_head.forward_train(x, img_metas, gt_semantic_seg, self.train_cfg) + losses.update(add_prefix(loss_aux, "aux")) return losses @@ -154,13 +140,11 @@ def forward_train(self, img, img_metas, gt_semantic_seg): losses = dict() - loss_decode = self._decode_head_forward_train(x, img_metas, - gt_semantic_seg) + loss_decode = self._decode_head_forward_train(x, img_metas, gt_semantic_seg) losses.update(loss_decode) if self.with_auxiliary_head: - loss_aux = self._auxiliary_head_forward_train( - x, img_metas, gt_semantic_seg) + loss_aux = self._auxiliary_head_forward_train(x, img_metas, gt_semantic_seg) losses.update(loss_aux) return losses @@ -191,24 +175,22 @@ def slide_inference(self, img, img_meta, rescale): x1 = max(x2 - w_crop, 0) crop_img = img[:, :, y1:y2, x1:x2] crop_seg_logit = self.encode_decode(crop_img, img_meta) - preds += F.pad(crop_seg_logit, - (int(x1), int(preds.shape[3] - x2), int(y1), - int(preds.shape[2] - y2))) + preds += F.pad(crop_seg_logit, (int(x1), int(preds.shape[3] - x2), int(y1), int(preds.shape[2] - y2))) count_mat[:, :, y1:y2, x1:x2] += 1 assert (count_mat == 0).sum() == 0 if torch.onnx.is_in_onnx_export(): # cast count_mat to constant while exporting to ONNX - count_mat = torch.from_numpy( - count_mat.cpu().detach().numpy()).to(device=img.device) + count_mat = torch.from_numpy(count_mat.cpu().detach().numpy()).to(device=img.device) preds = preds / count_mat if rescale: preds = resize( preds, - size=img_meta[0]['ori_shape'][:2], - mode='bilinear', + size=img_meta[0]["ori_shape"][:2], + mode="bilinear", align_corners=self.align_corners, - warning=False) + warning=False, + ) return preds def whole_inference(self, img, img_meta, rescale): @@ -220,13 +202,8 @@ def whole_inference(self, img, img_meta, rescale): if torch.onnx.is_in_onnx_export(): size = img.shape[2:] else: - size = img_meta[0]['ori_shape'][:2] - seg_logit = resize( - seg_logit, - size=size, - mode='bilinear', - align_corners=self.align_corners, - warning=False) + size = img_meta[0]["ori_shape"][:2] + seg_logit = resize(seg_logit, size=size, mode="bilinear", align_corners=self.align_corners, warning=False) return seg_logit @@ -246,22 +223,22 @@ def inference(self, img, img_meta, rescale): Tensor: The output segmentation map. """ - assert self.test_cfg.mode in ['slide', 'whole'] - ori_shape = img_meta[0]['ori_shape'] - assert all(_['ori_shape'] == ori_shape for _ in img_meta) - if self.test_cfg.mode == 'slide': + assert self.test_cfg.mode in ["slide", "whole"] + ori_shape = img_meta[0]["ori_shape"] + assert all(_["ori_shape"] == ori_shape for _ in img_meta) + if self.test_cfg.mode == "slide": seg_logit = self.slide_inference(img, img_meta, rescale) else: seg_logit = self.whole_inference(img, img_meta, rescale) output = F.softmax(seg_logit, dim=1) - flip = img_meta[0]['flip'] + flip = img_meta[0]["flip"] if flip: - flip_direction = img_meta[0]['flip_direction'] - assert flip_direction in ['horizontal', 'vertical'] - if flip_direction == 'horizontal': - output = output.flip(dims=(3, )) - elif flip_direction == 'vertical': - output = output.flip(dims=(2, )) + flip_direction = img_meta[0]["flip_direction"] + assert flip_direction in ["horizontal", "vertical"] + if flip_direction == "horizontal": + output = output.flip(dims=(3,)) + elif flip_direction == "vertical": + output = output.flip(dims=(2,)) return output diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/inverted_residual.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/inverted_residual.py index 331333e2..532ec832 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/inverted_residual.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/inverted_residual.py @@ -1,4 +1,4 @@ -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule from torch import nn from torch.utils import checkpoint as cp @@ -28,20 +28,21 @@ class InvertedResidual(nn.Module): Tensor: The output tensor. """ - def __init__(self, - in_channels, - out_channels, - stride, - expand_ratio, - dilation=1, - conv_cfg=None, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU6'), - with_cp=False): + def __init__( + self, + in_channels, + out_channels, + stride, + expand_ratio, + dilation=1, + conv_cfg=None, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU6"), + with_cp=False, + ): super(InvertedResidual, self).__init__() self.stride = stride - assert stride in [1, 2], f'stride must in [1, 2]. ' \ - f'But received {stride}.' + assert stride in [1, 2], f"stride must in [1, 2]. " f"But received {stride}." self.with_cp = with_cp self.use_res_connect = self.stride == 1 and in_channels == out_channels hidden_dim = int(round(in_channels * expand_ratio)) @@ -55,31 +56,36 @@ def __init__(self, kernel_size=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg)) - layers.extend([ - ConvModule( - in_channels=hidden_dim, - out_channels=hidden_dim, - kernel_size=3, - stride=stride, - padding=dilation, - dilation=dilation, - groups=hidden_dim, - conv_cfg=conv_cfg, - norm_cfg=norm_cfg, - act_cfg=act_cfg), - ConvModule( - in_channels=hidden_dim, - out_channels=out_channels, - kernel_size=1, - conv_cfg=conv_cfg, - norm_cfg=norm_cfg, - act_cfg=None) - ]) + act_cfg=act_cfg, + ) + ) + layers.extend( + [ + ConvModule( + in_channels=hidden_dim, + out_channels=hidden_dim, + kernel_size=3, + stride=stride, + padding=dilation, + dilation=dilation, + groups=hidden_dim, + conv_cfg=conv_cfg, + norm_cfg=norm_cfg, + act_cfg=act_cfg, + ), + ConvModule( + in_channels=hidden_dim, + out_channels=out_channels, + kernel_size=1, + conv_cfg=conv_cfg, + norm_cfg=norm_cfg, + act_cfg=None, + ), + ] + ) self.conv = nn.Sequential(*layers) def forward(self, x): - def _inner_forward(x): if self.use_res_connect: return x + self.conv(x) @@ -121,20 +127,22 @@ class InvertedResidualV3(nn.Module): Tensor: The output tensor. """ - def __init__(self, - in_channels, - out_channels, - mid_channels, - kernel_size=3, - stride=1, - se_cfg=None, - with_expand_conv=True, - conv_cfg=None, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU'), - with_cp=False): + def __init__( + self, + in_channels, + out_channels, + mid_channels, + kernel_size=3, + stride=1, + se_cfg=None, + with_expand_conv=True, + conv_cfg=None, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU"), + with_cp=False, + ): super(InvertedResidualV3, self).__init__() - self.with_res_shortcut = (stride == 1 and in_channels == out_channels) + self.with_res_shortcut = stride == 1 and in_channels == out_channels assert stride in [1, 2] self.with_cp = with_cp self.with_se = se_cfg is not None @@ -154,7 +162,8 @@ def __init__(self, padding=0, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) self.depthwise_conv = ConvModule( in_channels=mid_channels, out_channels=mid_channels, @@ -162,10 +171,10 @@ def __init__(self, stride=stride, padding=kernel_size // 2, groups=mid_channels, - conv_cfg=dict( - type='Conv2dAdaptivePadding') if stride == 2 else conv_cfg, + conv_cfg=dict(type="Conv2dAdaptivePadding") if stride == 2 else conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) if self.with_se: self.se = SELayer(**se_cfg) @@ -178,10 +187,10 @@ def __init__(self, padding=0, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=None) + act_cfg=None, + ) def forward(self, x): - def _inner_forward(x): out = x diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/res_layer.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/res_layer.py index e7979f9d..765e40be 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/res_layer.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/res_layer.py @@ -1,4 +1,4 @@ -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import build_conv_layer, build_norm_layer +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import build_conv_layer, build_norm_layer from torch import nn as nn @@ -23,19 +23,21 @@ class ResLayer(nn.Sequential): Default: False """ - def __init__(self, - block, - inplanes, - planes, - num_blocks, - stride=1, - dilation=1, - avg_down=False, - conv_cfg=None, - norm_cfg=dict(type='BN'), - multi_grid=None, - contract_dilation=False, - **kwargs): + def __init__( + self, + block, + inplanes, + planes, + num_blocks, + stride=1, + dilation=1, + avg_down=False, + conv_cfg=None, + norm_cfg=dict(type="BN"), + multi_grid=None, + contract_dilation=False, + **kwargs, + ): self.block = block downsample = None @@ -45,21 +47,16 @@ def __init__(self, if avg_down: conv_stride = 1 downsample.append( - nn.AvgPool2d( - kernel_size=stride, - stride=stride, - ceil_mode=True, - count_include_pad=False)) - downsample.extend([ - build_conv_layer( - conv_cfg, - inplanes, - planes * block.expansion, - kernel_size=1, - stride=conv_stride, - bias=False), - build_norm_layer(norm_cfg, planes * block.expansion)[1] - ]) + nn.AvgPool2d(kernel_size=stride, stride=stride, ceil_mode=True, count_include_pad=False) + ) + downsample.extend( + [ + build_conv_layer( + conv_cfg, inplanes, planes * block.expansion, kernel_size=1, stride=conv_stride, bias=False + ), + build_norm_layer(norm_cfg, planes * block.expansion)[1], + ] + ) downsample = nn.Sequential(*downsample) layers = [] @@ -79,7 +76,9 @@ def __init__(self, downsample=downsample, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - **kwargs)) + **kwargs, + ) + ) inplanes = planes * block.expansion for i in range(1, num_blocks): layers.append( @@ -90,5 +89,7 @@ def __init__(self, dilation=dilation if multi_grid is None else multi_grid[i], conv_cfg=conv_cfg, norm_cfg=norm_cfg, - **kwargs)) + **kwargs, + ) + ) super(ResLayer, self).__init__(*layers) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/se_layer.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/se_layer.py index a56fd04f..faadb9b1 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/se_layer.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/se_layer.py @@ -1,6 +1,6 @@ -import comfy_controlnet_preprocessors.uniformer.mmcv as mmcv +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv as mmcv import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule from .make_divisible import make_divisible @@ -23,12 +23,13 @@ class SELayer(nn.Module): divisor=6.0)). """ - def __init__(self, - channels, - ratio=16, - conv_cfg=None, - act_cfg=(dict(type='ReLU'), - dict(type='HSigmoid', bias=3.0, divisor=6.0))): + def __init__( + self, + channels, + ratio=16, + conv_cfg=None, + act_cfg=(dict(type="ReLU"), dict(type="HSigmoid", bias=3.0, divisor=6.0)), + ): super(SELayer, self).__init__() if isinstance(act_cfg, dict): act_cfg = (act_cfg, act_cfg) @@ -41,14 +42,16 @@ def __init__(self, kernel_size=1, stride=1, conv_cfg=conv_cfg, - act_cfg=act_cfg[0]) + act_cfg=act_cfg[0], + ) self.conv2 = ConvModule( in_channels=make_divisible(channels // ratio, 8), out_channels=channels, kernel_size=1, stride=1, conv_cfg=conv_cfg, - act_cfg=act_cfg[1]) + act_cfg=act_cfg[1], + ) def forward(self, x): out = self.global_avgpool(x) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/self_attention_block.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/self_attention_block.py index f8522680..714433c4 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/self_attention_block.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/self_attention_block.py @@ -1,5 +1,5 @@ import torch -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, constant_init +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, constant_init from torch import nn as nn from torch.nn import functional as F @@ -29,11 +29,25 @@ class SelfAttentionBlock(nn.Module): act_cfg (dict|None): Config of activation layers. """ - def __init__(self, key_in_channels, query_in_channels, channels, - out_channels, share_key_query, query_downsample, - key_downsample, key_query_num_convs, value_out_num_convs, - key_query_norm, value_out_norm, matmul_norm, with_out, - conv_cfg, norm_cfg, act_cfg): + def __init__( + self, + key_in_channels, + query_in_channels, + channels, + out_channels, + share_key_query, + query_downsample, + key_downsample, + key_query_num_convs, + value_out_num_convs, + key_query_norm, + value_out_norm, + matmul_norm, + with_out, + conv_cfg, + norm_cfg, + act_cfg, + ): super(SelfAttentionBlock, self).__init__() if share_key_query: assert key_in_channels == query_in_channels @@ -52,7 +66,8 @@ def __init__(self, key_in_channels, query_in_channels, channels, use_conv_module=key_query_norm, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) if share_key_query: self.query_project = self.key_project else: @@ -63,7 +78,8 @@ def __init__(self, key_in_channels, query_in_channels, channels, use_conv_module=key_query_norm, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) self.value_project = self.build_project( key_in_channels, channels if with_out else out_channels, @@ -71,7 +87,8 @@ def __init__(self, key_in_channels, query_in_channels, channels, use_conv_module=value_out_norm, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) if with_out: self.out_project = self.build_project( channels, @@ -80,7 +97,8 @@ def __init__(self, key_in_channels, query_in_channels, channels, use_conv_module=value_out_norm, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) else: self.out_project = None @@ -96,28 +114,12 @@ def init_weights(self): if not isinstance(self.out_project, ConvModule): constant_init(self.out_project, 0) - def build_project(self, in_channels, channels, num_convs, use_conv_module, - conv_cfg, norm_cfg, act_cfg): + def build_project(self, in_channels, channels, num_convs, use_conv_module, conv_cfg, norm_cfg, act_cfg): """Build projection layer for key/query/value/out.""" if use_conv_module: - convs = [ - ConvModule( - in_channels, - channels, - 1, - conv_cfg=conv_cfg, - norm_cfg=norm_cfg, - act_cfg=act_cfg) - ] + convs = [ConvModule(in_channels, channels, 1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg)] for _ in range(num_convs - 1): - convs.append( - ConvModule( - channels, - channels, - 1, - conv_cfg=conv_cfg, - norm_cfg=norm_cfg, - act_cfg=act_cfg)) + convs.append(ConvModule(channels, channels, 1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg)) else: convs = [nn.Conv2d(in_channels, channels, 1)] for _ in range(num_convs - 1): @@ -148,7 +150,7 @@ def forward(self, query_feats, key_feats): sim_map = torch.matmul(query, key) if self.matmul_norm: - sim_map = (self.channels**-.5) * sim_map + sim_map = (self.channels**-0.5) * sim_map sim_map = F.softmax(sim_map, dim=-1) context = torch.matmul(sim_map, value) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/up_conv_block.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/up_conv_block.py index ceac2fcc..fa0a17da 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/up_conv_block.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/models/utils/up_conv_block.py @@ -1,6 +1,6 @@ import torch import torch.nn as nn -from comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, build_upsample_layer +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.cnn import ConvModule, build_upsample_layer class UpConvBlock(nn.Module): @@ -41,24 +41,26 @@ class UpConvBlock(nn.Module): plugins (dict): plugins for convolutional layers. Default: None. """ - def __init__(self, - conv_block, - in_channels, - skip_channels, - out_channels, - num_convs=2, - stride=1, - dilation=1, - with_cp=False, - conv_cfg=None, - norm_cfg=dict(type='BN'), - act_cfg=dict(type='ReLU'), - upsample_cfg=dict(type='InterpConv'), - dcn=None, - plugins=None): + def __init__( + self, + conv_block, + in_channels, + skip_channels, + out_channels, + num_convs=2, + stride=1, + dilation=1, + with_cp=False, + conv_cfg=None, + norm_cfg=dict(type="BN"), + act_cfg=dict(type="ReLU"), + upsample_cfg=dict(type="InterpConv"), + dcn=None, + plugins=None, + ): super(UpConvBlock, self).__init__() - assert dcn is None, 'Not implemented yet.' - assert plugins is None, 'Not implemented yet.' + assert dcn is None, "Not implemented yet." + assert plugins is None, "Not implemented yet." self.conv_block = conv_block( in_channels=2 * skip_channels, @@ -71,7 +73,8 @@ def __init__(self, norm_cfg=norm_cfg, act_cfg=act_cfg, dcn=None, - plugins=None) + plugins=None, + ) if upsample_cfg is not None: self.upsample = build_upsample_layer( cfg=upsample_cfg, @@ -79,7 +82,8 @@ def __init__(self, out_channels=skip_channels, with_cp=with_cp, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) else: self.upsample = ConvModule( in_channels, @@ -89,7 +93,8 @@ def __init__(self, padding=0, conv_cfg=conv_cfg, norm_cfg=norm_cfg, - act_cfg=act_cfg) + act_cfg=act_cfg, + ) def forward(self, skip, x): """Forward function.""" diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/utils/collect_env.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/utils/collect_env.py index f002e0aa..abed6c55 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/utils/collect_env.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/utils/collect_env.py @@ -1,17 +1,17 @@ -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import collect_env as collect_base_env -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import get_git_hash +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import collect_env as collect_base_env +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import get_git_hash -import comfy_controlnet_preprocessors.uniformer.mmseg as mmseg +import hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmseg as mmseg def collect_env(): """Collect the information of the running environments.""" env_info = collect_base_env() - env_info['MMSegmentation'] = f'{mmseg.__version__}+{get_git_hash()[:7]}' + env_info["MMSegmentation"] = f"{mmseg.__version__}+{get_git_hash()[:7]}" return env_info -if __name__ == '__main__': +if __name__ == "__main__": for name, val in collect_env().items(): - print('{}: {}'.format(name, val)) + print("{}: {}".format(name, val)) diff --git a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/utils/logger.py b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/utils/logger.py index 974b1292..4d59016a 100644 --- a/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/utils/logger.py +++ b/hordelib/nodes/comfy_controlnet_preprocessors/uniformer/mmseg/utils/logger.py @@ -1,6 +1,6 @@ import logging -from comfy_controlnet_preprocessors.uniformer.mmcv.utils import get_logger +from hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv.utils import get_logger def get_root_logger(log_file=None, log_level=logging.INFO): @@ -22,6 +22,6 @@ def get_root_logger(log_file=None, log_level=logging.INFO): logging.Logger: The root logger. """ - logger = get_logger(name='mmseg', log_file=log_file, log_level=log_level) + logger = get_logger(name="mmseg", log_file=log_file, log_level=log_level) return logger diff --git a/hordelib/nodes/facerestore/__init__.py b/hordelib/nodes/facerestore/__init__.py index 4872d2fc..b2fdac14 100644 --- a/hordelib/nodes/facerestore/__init__.py +++ b/hordelib/nodes/facerestore/__init__.py @@ -5,8 +5,8 @@ import numpy as np import cv2 import math -from facerestore.facelib.utils.face_restoration_helper import FaceRestoreHelper -from facerestore.facelib.detection.retinaface import retinaface +from hordelib.nodes.facerestore.facelib.utils.face_restoration_helper import FaceRestoreHelper +from hordelib.nodes.facerestore.facelib.detection.retinaface import retinaface from torchvision.transforms.functional import normalize import threading from loguru import logger diff --git a/hordelib/nodes/facerestore/facelib/detection/__init__.py b/hordelib/nodes/facerestore/facelib/detection/__init__.py index a0dc42f5..7ec791bf 100644 --- a/hordelib/nodes/facerestore/facelib/detection/__init__.py +++ b/hordelib/nodes/facerestore/facelib/detection/__init__.py @@ -3,9 +3,9 @@ from torch import nn from copy import deepcopy -from facerestore.facelib.utils import load_file_from_url -from facerestore.facelib.utils import download_pretrained_models -from facerestore.facelib.detection.yolov5face.models.common import Conv +from hordelib.nodes.facerestore.facelib.utils import load_file_from_url +from hordelib.nodes.facerestore.facelib.utils import download_pretrained_models +from hordelib.nodes.facerestore.facelib.detection.yolov5face.models.common import Conv from .retinaface.retinaface import RetinaFace from .yolov5face.face_detector import YoloDetector diff --git a/hordelib/nodes/facerestore/facelib/detection/retinaface/retinaface.py b/hordelib/nodes/facerestore/facelib/detection/retinaface/retinaface.py index 94dda7a4..16881365 100644 --- a/hordelib/nodes/facerestore/facelib/detection/retinaface/retinaface.py +++ b/hordelib/nodes/facerestore/facelib/detection/retinaface/retinaface.py @@ -6,11 +6,11 @@ from PIL import Image from torchvision.models._utils import IntermediateLayerGetter as IntermediateLayerGetter -from facerestore.facelib.detection.align_trans import ( +from hordelib.nodes.facerestore.facelib.detection.align_trans import ( get_reference_facial_points, warp_and_crop_face, ) -from facerestore.facelib.detection.retinaface.retinaface_net import ( +from hordelib.nodes.facerestore.facelib.detection.retinaface.retinaface_net import ( FPN, SSH, MobileNetV1, @@ -18,7 +18,7 @@ make_class_head, make_landmark_head, ) -from facerestore.facelib.detection.retinaface.retinaface_utils import ( +from hordelib.nodes.facerestore.facelib.detection.retinaface.retinaface_utils import ( PriorBox, batched_decode, batched_decode_landm, diff --git a/hordelib/nodes/facerestore/facelib/detection/yolov5face/face_detector.py b/hordelib/nodes/facerestore/facelib/detection/yolov5face/face_detector.py index 4f62ad76..5ea44d65 100644 --- a/hordelib/nodes/facerestore/facelib/detection/yolov5face/face_detector.py +++ b/hordelib/nodes/facerestore/facelib/detection/yolov5face/face_detector.py @@ -7,10 +7,10 @@ import torch from torch import nn -from facerestore.facelib.detection.yolov5face.models.common import Conv -from facerestore.facelib.detection.yolov5face.models.yolo import Model -from facerestore.facelib.detection.yolov5face.utils.datasets import letterbox -from facerestore.facelib.detection.yolov5face.utils.general import ( +from hordelib.nodes.facerestore.facelib.detection.yolov5face.models.common import Conv +from hordelib.nodes.facerestore.facelib.detection.yolov5face.models.yolo import Model +from hordelib.nodes.facerestore.facelib.detection.yolov5face.utils.datasets import letterbox +from hordelib.nodes.facerestore.facelib.detection.yolov5face.utils.general import ( check_img_size, non_max_suppression_face, scale_coords, diff --git a/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/common.py b/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/common.py index 199532aa..9c6c5eaa 100644 --- a/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/common.py +++ b/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/common.py @@ -6,8 +6,8 @@ import torch from torch import nn -from facerestore.facelib.detection.yolov5face.utils.datasets import letterbox -from facerestore.facelib.detection.yolov5face.utils.general import ( +from hordelib.nodes.facerestore.facelib.detection.yolov5face.utils.datasets import letterbox +from hordelib.nodes.facerestore.facelib.detection.yolov5face.utils.general import ( make_divisible, non_max_suppression, scale_coords, @@ -41,17 +41,11 @@ def DWConv(c1, c2, k=1, s=1, act=True): class Conv(nn.Module): # Standard convolution - def __init__( - self, c1, c2, k=1, s=1, p=None, g=1, act=True - ): # ch_in, ch_out, kernel, stride, padding, groups + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups super().__init__() self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) self.bn = nn.BatchNorm2d(c2) - self.act = ( - nn.SiLU() - if act is True - else (act if isinstance(act, nn.Module) else nn.Identity()) - ) + self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity()) def forward(self, x): return self.act(self.bn(self.conv(x))) @@ -79,9 +73,7 @@ def forward(self, x): class Bottleneck(nn.Module): # Standard bottleneck - def __init__( - self, c1, c2, shortcut=True, g=1, e=0.5 - ): # ch_in, ch_out, shortcut, groups, expansion + def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion super().__init__() c_ = int(c2 * e) # hidden channels self.cv1 = Conv(c1, c_, 1, 1) @@ -94,9 +86,7 @@ def forward(self, x): class BottleneckCSP(nn.Module): # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks - def __init__( - self, c1, c2, n=1, shortcut=True, g=1, e=0.5 - ): # ch_in, ch_out, number, shortcut, groups, expansion + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion super().__init__() c_ = int(c2 * e) # hidden channels self.cv1 = Conv(c1, c_, 1, 1) @@ -105,9 +95,7 @@ def __init__( self.cv4 = Conv(2 * c_, c2, 1, 1) self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) self.act = nn.LeakyReLU(0.1, inplace=True) - self.m = nn.Sequential( - *(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)) - ) + self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n))) def forward(self, x): y1 = self.cv3(self.m(self.cv1(x))) @@ -117,17 +105,13 @@ def forward(self, x): class C3(nn.Module): # CSP Bottleneck with 3 convolutions - def __init__( - self, c1, c2, n=1, shortcut=True, g=1, e=0.5 - ): # ch_in, ch_out, number, shortcut, groups, expansion + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion super().__init__() c_ = int(c2 * e) # hidden channels self.cv1 = Conv(c1, c_, 1, 1) self.cv2 = Conv(c1, c_, 1, 1) self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2) - self.m = nn.Sequential( - *(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)) - ) + self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n))) def forward(self, x): return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1)) @@ -145,13 +129,9 @@ def __init__(self, inp, oup, stride): if self.stride > 1: self.branch1 = nn.Sequential( - self.depthwise_conv( - inp, inp, kernel_size=3, stride=self.stride, padding=1 - ), + self.depthwise_conv(inp, inp, kernel_size=3, stride=self.stride, padding=1), nn.BatchNorm2d(inp), - nn.Conv2d( - inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False - ), + nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(branch_features), nn.SiLU(), ) @@ -210,9 +190,7 @@ def __init__(self, c1, c2, k=(5, 9, 13)): c_ = c1 // 2 # hidden channels self.cv1 = Conv(c1, c_, 1, 1) self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) - self.m = nn.ModuleList( - [nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k] - ) + self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) def forward(self, x): x = self.cv1(x) @@ -221,9 +199,7 @@ def forward(self, x): class Focus(nn.Module): # Focus wh information into c-space - def __init__( - self, c1, c2, k=1, s=1, p=None, g=1, act=True - ): # ch_in, ch_out, kernel, stride, padding, groups + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups super().__init__() self.conv = Conv(c1 * 4, c2, k, s, p, g, act) @@ -258,9 +234,7 @@ class NMS(nn.Module): classes = None # (optional list) filter by class def forward(self, x): - return non_max_suppression( - x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes - ) + return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) class AutoShape(nn.Module): @@ -275,9 +249,7 @@ def __init__(self, model): self.model = model.eval() def autoshape(self): - print( - "autoShape already enabled, skipping... " - ) # model already converted to model.autoshape() + print("autoShape already enabled, skipping... ") # model already converted to model.autoshape() return self def forward(self, imgs, size=640, augment=False, profile=False): @@ -290,31 +262,22 @@ def forward(self, imgs, size=640, augment=False, profile=False): p = next(self.model.parameters()) # for device and type if isinstance(imgs, torch.Tensor): # torch - return self.model( - imgs.to(p.device).type_as(p), augment, profile - ) # inference + return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference # Pre-process - n, imgs = ( - (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) - ) # number of images, list of images + n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images shape0, shape1 = [], [] # image and inference shapes for i, im in enumerate(imgs): im = np.array(im) # to numpy if im.shape[0] < 5: # image in CHW im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1) - im = ( - im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) - ) # enforce 3ch input + im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input s = im.shape[:2] # HWC shape0.append(s) # image shape g = size / max(s) # gain shape1.append([y * g for y in s]) imgs[i] = im # update - shape1 = [ - make_divisible(x, int(self.stride.max())) - for x in np.stack(shape1, 0).max(0) - ] # inference shape + shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad x = np.stack(x, 0) if n > 1 else x[0][None] # stack x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW @@ -323,9 +286,7 @@ def forward(self, imgs, size=640, augment=False, profile=False): # Inference with torch.no_grad(): y = self.model(x, augment, profile)[0] # forward - y = non_max_suppression( - y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes - ) # NMS + y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS # Post-process for i in range(n): @@ -340,8 +301,7 @@ def __init__(self, imgs, pred, names=None): super().__init__() d = pred[0].device # device gn = [ - torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1.0, 1.0], device=d) - for im in imgs + torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1.0, 1.0], device=d) for im in imgs ] # normalizations self.imgs = imgs # list of images as numpy arrays self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls) @@ -357,10 +317,7 @@ def __len__(self): def tolist(self): # return a list of Detections objects, i.e. 'for result in results.tolist():' - x = [ - Detections([self.imgs[i]], [self.pred[i]], self.names) - for i in range(self.n) - ] + x = [Detections([self.imgs[i]], [self.pred[i]], self.names) for i in range(self.n)] for d in x: for k in ["imgs", "pred", "xyxy", "xyxyn", "xywh", "xywhn"]: setattr(d, k, getattr(d, k)[0]) # pop out of list diff --git a/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/experimental.py b/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/experimental.py index db7c4eb2..c33acba5 100644 --- a/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/experimental.py +++ b/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/experimental.py @@ -4,7 +4,7 @@ import torch from torch import nn -from facerestore.facelib.detection.yolov5face.models.common import Conv +from hordelib.nodes.facerestore.facelib.detection.yolov5face.models.common import Conv class CrossConv(nn.Module): diff --git a/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/yolo.py b/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/yolo.py index abca50d7..b11af3e4 100644 --- a/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/yolo.py +++ b/hordelib/nodes/facerestore/facelib/detection/yolov5face/models/yolo.py @@ -6,7 +6,7 @@ import yaml # for torch hub from torch import nn -from facerestore.facelib.detection.yolov5face.models.common import ( +from hordelib.nodes.facerestore.facelib.detection.yolov5face.models.common import ( C3, NMS, SPP, @@ -20,13 +20,13 @@ ShuffleV2Block, StemBlock, ) -from facerestore.facelib.detection.yolov5face.models.experimental import ( +from hordelib.nodes.facerestore.facelib.detection.yolov5face.models.experimental import ( CrossConv, MixConv2d, ) -from facerestore.facelib.detection.yolov5face.utils.autoanchor import check_anchor_order -from facerestore.facelib.detection.yolov5face.utils.general import make_divisible -from facerestore.facelib.detection.yolov5face.utils.torch_utils import ( +from hordelib.nodes.facerestore.facelib.detection.yolov5face.utils.autoanchor import check_anchor_order +from hordelib.nodes.facerestore.facelib.detection.yolov5face.utils.general import make_divisible +from hordelib.nodes.facerestore.facelib.detection.yolov5face.utils.torch_utils import ( copy_attr, fuse_conv_and_bn, ) diff --git a/hordelib/nodes/facerestore/facelib/parsing/__init__.py b/hordelib/nodes/facerestore/facelib/parsing/__init__.py index 89e14c94..b95f2d8b 100644 --- a/hordelib/nodes/facerestore/facelib/parsing/__init__.py +++ b/hordelib/nodes/facerestore/facelib/parsing/__init__.py @@ -1,6 +1,6 @@ import torch -from facerestore.facelib.utils import load_file_from_url +from hordelib.nodes.facerestore.facelib.utils import load_file_from_url from .bisenet import BiSeNet from .parsenet import ParseNet diff --git a/hordelib/nodes/facerestore/facelib/utils/face_restoration_helper.py b/hordelib/nodes/facerestore/facelib/utils/face_restoration_helper.py index 2aa24d86..7917344a 100644 --- a/hordelib/nodes/facerestore/facelib/utils/face_restoration_helper.py +++ b/hordelib/nodes/facerestore/facelib/utils/face_restoration_helper.py @@ -4,9 +4,9 @@ import torch from torchvision.transforms.functional import normalize -from facerestore.facelib.detection import init_detection_model -from facerestore.facelib.parsing import init_parsing_model -from facerestore.facelib.utils.misc import img2tensor, imwrite +from hordelib.nodes.facerestore.facelib.detection import init_detection_model +from hordelib.nodes.facerestore.facelib.parsing import init_parsing_model +from hordelib.nodes.facerestore.facelib.utils.misc import img2tensor, imwrite def get_largest_face(det_faces, h, w): diff --git a/hordelib/nodes/facerestore/facelib/utils/face_utils.py b/hordelib/nodes/facerestore/facelib/utils/face_utils.py index 06e6f4fb..5ee39570 100644 --- a/hordelib/nodes/facerestore/facelib/utils/face_utils.py +++ b/hordelib/nodes/facerestore/facelib/utils/face_utils.py @@ -245,8 +245,8 @@ def paste_face_back(img, face, inverse_affine): if __name__ == "__main__": import os - from facerestore.facelib.detection import init_detection_model - from facerestore.facelib.utils.face_restoration_helper import get_largest_face + from hordelib.nodes.facerestore.facelib.detection import init_detection_model + from hordelib.nodes.facerestore.facelib.utils.face_restoration_helper import get_largest_face img_path = "/home/wxt/datasets/ffhq/ffhq_wild/00009.png" img_name = os.path.splitext(os.path.basename(img_path))[0] diff --git a/hordelib/pyinstaller_hooks/__init__.py b/hordelib/pyinstaller_hooks/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/hordelib/pyinstaller_hooks/hook-hordelib.horde.py b/hordelib/pyinstaller_hooks/hook-hordelib.horde.py new file mode 100644 index 00000000..c565c03d --- /dev/null +++ b/hordelib/pyinstaller_hooks/hook-hordelib.horde.py @@ -0,0 +1,53 @@ +import pathlib + +from PyInstaller.utils.hooks import collect_data_files # type: ignore + +hiddenimports = [ + "torch", + "torchvision", + "torchsde", + "transformers", + "safetensors", + "scipy", + "tqdm", + "psutil", + "opencv-python", + "opencv-contrib-python", + "timm", + "addict", + "fairscale", + "scikit-image", + "mediapipe", + "unidecode", + "fuzzywuzzy", + "open_clip", + "clip-interrogator", +] +warn_on_missing_hiddenimports = True + +# module_collection_mode = "pyz+py" + +project_root = pathlib.Path(__file__).parent.parent + +datas = [ + (str(project_root / "_version.py"), "hordelib"), +] + +excluded_folders = [ + "__pycache__", + ".git", +] + +# Add every file in the `hordelib/_comfyui` directory to datas, copying the directory structure +comfyui_dir = "hordelib._comfyui" +datas += collect_data_files(comfyui_dir, "hordelib/_comfyui", excludes=excluded_folders) + +nodes_dir = "hordelib.nodes" +datas += collect_data_files(nodes_dir, "hordelib/nodes") + + +pipelines_dir = "hordelib.pipelines" +pipeline_designs_dir = "hordelib.pipeline_designs" + +datas += collect_data_files(pipelines_dir, "hordelib/pipelines") +datas += collect_data_files(pipeline_designs_dir, "hordelib/pipeline_designs") diff --git a/hordelib/utils/distance.py b/hordelib/utils/distance.py index 41288bf1..2ff599dc 100644 --- a/hordelib/utils/distance.py +++ b/hordelib/utils/distance.py @@ -231,11 +231,11 @@ def cv2_image_similarity( raise ValueError("Images must be of same size.") # Convert the images to grayscale - img1 = cv2.cvtColor(np.array(img1), cv2.COLOR_BGR2GRAY) - img2 = cv2.cvtColor(np.array(img2), cv2.COLOR_BGR2GRAY) + img1 = cv2.cvtColor(np.array(img1), cv2.COLOR_BGR2GRAY) # type: ignore # FIXME + img2 = cv2.cvtColor(np.array(img2), cv2.COLOR_BGR2GRAY) # type: ignore # FIXME # Calculate the similarity between the two images - similarity = cv2.matchTemplate(img1, img2, cv2.TM_CCOEFF_NORMED)[0][0] + similarity = cv2.matchTemplate(img1, img2, cv2.TM_CCOEFF_NORMED)[0][0] # type: ignore # FIXME # Get the primitive (float) value of the similarity similarity = similarity.item() diff --git a/mypy.ini b/mypy.ini index 79f4c2ec..a0a5b2e5 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1,5 +1,5 @@ [mypy] -exclude = (build|ComfyUI|comfy_controlnet_preprocessors|facerestore|comfy_horde\.py|examples|diffusers) +exclude = (build|dist|ComfyUI|comfy_controlnet_preprocessors|facerestore|comfy_horde\.py|examples|diffusers) [mypy-hordelib.nodes.comfy_controlnet_preprocessors.*] diff --git a/pyinstaller.py b/pyinstaller.py new file mode 100644 index 00000000..447514f7 --- /dev/null +++ b/pyinstaller.py @@ -0,0 +1,21 @@ +from pathlib import Path + +import PyInstaller.__main__ # type: ignore + +HERE = Path(__file__).parent.absolute() +entry_point = str(HERE / "pyinstaller_test_entrypoint.py") + + +def install(): + PyInstaller.__main__.run( + [ + "--clean", + "--log-level=DEBUG", + "--additional-hooks-dir", + str(HERE / "pyinstaller_hooks"), + entry_point, + ], + ) + + +install() diff --git a/pyinstaller_test_entrypoint.py b/pyinstaller_test_entrypoint.py new file mode 100644 index 00000000..58ec849a --- /dev/null +++ b/pyinstaller_test_entrypoint.py @@ -0,0 +1,36 @@ +from hordelib.horde import HordeLib +from hordelib.initialisation import initialise + + +def main(): + initialise() + + hordelib_instance = HordeLib() + + data = { + "sampler_name": "k_dpmpp_2m", + "cfg_scale": 7.5, + "denoising_strength": 1.0, + "seed": 123456789, + "height": 512.1, # test param fix + "width": 512.1, # test param fix + "karras": False, + "tiling": False, + "hires_fix": False, + "clip_skip": 1, + "control_type": None, + "image_is_control": False, + "return_control_map": False, + "prompt": "an ancient llamia monster", + "ddim_steps": 25, + "n_iter": 1, + "model": "Deliberate", + } + pil_image = hordelib_instance.basic_inference_single_image(data).image + + img_filename = "text_to_image.png" + pil_image.save(f"images/{img_filename}", quality=100) + + +if __name__ == "__main__": + main() From 5473325a774632e9626c7cbba150309e4daa7d7f Mon Sep 17 00:00:00 2001 From: tazlin Date: Tue, 6 Feb 2024 10:40:27 -0500 Subject: [PATCH 4/6] build: include pyinstaller entry points --- setup.cfg | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 setup.cfg diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 00000000..05350b8c --- /dev/null +++ b/setup.cfg @@ -0,0 +1,4 @@ +[options.entry_points] +pyinstaller40 = + hook-dirs = hordelib.__pyinstaller:get_hook_dirs + tests = hordelib.__pyinstaller:get_PyInstaller_tests From c161575d8be5a2517414db7b6da1238fb3509e47 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 6 Feb 2024 13:55:50 +0000 Subject: [PATCH 5/6] build(deps-dev): bump the python-packages group with 4 updates Updates the requirements on [pytest](https://github.com/pytest-dev/pytest), [black](https://github.com/psf/black), [ruff](https://github.com/astral-sh/ruff) and [tox](https://github.com/tox-dev/tox) to permit the latest version. Updates `pytest` from 7.4.4 to 8.0.0 - [Release notes](https://github.com/pytest-dev/pytest/releases) - [Changelog](https://github.com/pytest-dev/pytest/blob/main/CHANGELOG.rst) - [Commits](https://github.com/pytest-dev/pytest/compare/7.4.4...8.0.0) Updates `black` from 23.12.1 to 24.1.1 - [Release notes](https://github.com/psf/black/releases) - [Changelog](https://github.com/psf/black/blob/main/CHANGES.md) - [Commits](https://github.com/psf/black/compare/23.12.1...24.1.1) Updates `ruff` from 0.1.11 to 0.2.1 - [Release notes](https://github.com/astral-sh/ruff/releases) - [Changelog](https://github.com/astral-sh/ruff/blob/main/CHANGELOG.md) - [Commits](https://github.com/astral-sh/ruff/compare/v0.1.11...v0.2.1) Updates `tox` to 4.12.1 - [Release notes](https://github.com/tox-dev/tox/releases) - [Changelog](https://github.com/tox-dev/tox/blob/main/docs/changelog.rst) - [Commits](https://github.com/tox-dev/tox/compare/4.11.4...4.12.1) --- updated-dependencies: - dependency-name: pytest dependency-type: direct:development update-type: version-update:semver-major dependency-group: python-packages - dependency-name: black dependency-type: direct:development update-type: version-update:semver-major dependency-group: python-packages - dependency-name: ruff dependency-type: direct:development update-type: version-update:semver-minor dependency-group: python-packages - dependency-name: tox dependency-type: direct:development dependency-group: python-packages ... Signed-off-by: dependabot[bot] --- requirements.dev.txt | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/requirements.dev.txt b/requirements.dev.txt index eb697e2e..c488dde3 100644 --- a/requirements.dev.txt +++ b/requirements.dev.txt @@ -1,8 +1,8 @@ -pytest==7.4.4 +pytest==8.0.0 mypy==1.8.0 -black==23.12.1 -ruff==0.1.11 -tox~=4.11.4 +black==24.1.1 +ruff==0.2.1 +tox~=4.12.1 pre-commit~=3.6.0 build>=0.10.0 coverage>=7.2.7 From 1fd63a55fcdfd366d21920e3fdd4e5b434e2fe64 Mon Sep 17 00:00:00 2001 From: tazlin Date: Tue, 6 Feb 2024 12:27:02 -0500 Subject: [PATCH 6/6] chore: update pre-commit hook versions also adds `horde_sdk` to mypy hook addtl. deps, and applies the `black` format changes now required --- .pre-commit-config.yaml | 6 +++--- hordelib/model_manager/base.py | 25 ++++++++++++++----------- hordelib/model_manager/hyper.py | 7 ++++--- 3 files changed, 21 insertions(+), 17 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 91a017e4..d8932645 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -6,12 +6,12 @@ repos: - id: end-of-file-fixer - id: trailing-whitespace - repo: https://github.com/psf/black - rev: 23.12.1 + rev: 24.1.1 hooks: - id: black exclude: ^hordelib/nodes/.*\..*$ - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.1.11 + rev: v0.2.1 hooks: - id: ruff - repo: https://github.com/pre-commit/mirrors-mypy @@ -19,4 +19,4 @@ repos: hooks: - id: mypy exclude: ^examples/.*$ # FIXME - additional_dependencies: [pydantic, strenum, types-colorama, types-docutils, types-Pillow, types-psutil, types-Pygments, types-pywin32, types-PyYAML, types-regex, types-requests, types-setuptools, types-tabulate, types-tqdm, types-urllib3] + additional_dependencies: [pydantic, strenum, types-colorama, types-docutils, types-Pillow, types-psutil, types-Pygments, types-pywin32, types-PyYAML, types-regex, types-requests, types-setuptools, types-tabulate, types-tqdm, types-urllib3, horde_sdk] diff --git a/hordelib/model_manager/base.py b/hordelib/model_manager/base.py index 8ba81486..1ef6022a 100644 --- a/hordelib/model_manager/base.py +++ b/hordelib/model_manager/base.py @@ -498,17 +498,20 @@ def download_file( response.raise_for_status() # Write the content to file in chunks - with open(partial_pathname, "ab") as f, tqdm( - # all optional kwargs - unit="B", - initial=partial_size, - unit_scale=True, - unit_divisor=1024, - miniters=1, - desc=filename, - total=remote_file_size + partial_size, - # disable=UserSettings.download_progress_callback is not None, - ) as pbar: + with ( + open(partial_pathname, "ab") as f, + tqdm( + # all optional kwargs + unit="B", + initial=partial_size, + unit_scale=True, + unit_divisor=1024, + miniters=1, + desc=filename, + total=remote_file_size + partial_size, + # disable=UserSettings.download_progress_callback is not None, + ) as pbar, + ): downloaded = partial_size for chunk in response.iter_content(chunk_size=1024 * 1024 * 16): response.raise_for_status() diff --git a/hordelib/model_manager/hyper.py b/hordelib/model_manager/hyper.py index 3c8bfa5e..93bb7c7e 100644 --- a/hordelib/model_manager/hyper.py +++ b/hordelib/model_manager/hyper.py @@ -1,4 +1,5 @@ """Home for the controller class ModelManager, and related meta information.""" + import os import threading from collections.abc import Callable, Iterable @@ -180,9 +181,9 @@ def init_model_managers( self.active_model_managers.append( resolve_manager_to_load_type( multiprocessing_lock=multiprocessing_lock, - civitai_api_token=os.environ["CIVIT_API_TOKEN"] - if "CIVIT_API_TOKEN" in os.environ.keys() - else None, + civitai_api_token=( + os.environ["CIVIT_API_TOKEN"] if "CIVIT_API_TOKEN" in os.environ.keys() else None + ), ), )