add built-in extension system

add support for adding upscalers in extensions
move LDSR, ScuNET and SwinIR to built-in extensions

13 files changed, 42 insertions, 2710 deletions

diff --git a/modules/devices.py b/modules/devices.py
index d6a76844..f8cffae1 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -44,6 +44,15 @@ def get_optimal_device():
     return cpu
 
 
+def get_device_for(task):
+    from modules import shared
+
+    if task in shared.cmd_opts.use_cpu:
+        return cpu
+
+    return get_optimal_device()
+
+
 def torch_gc():
     if torch.cuda.is_available():
         with torch.cuda.device(get_cuda_device_string()):
@@ -67,7 +76,7 @@ def enable_tf32():
 errors.run(enable_tf32, "Enabling TF32")
 
 cpu = torch.device("cpu")
-device = device_interrogate = device_gfpgan = device_swinir = device_esrgan = device_scunet = device_codeformer = None
+device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = None
 dtype = torch.float16
 dtype_vae = torch.float16
 
diff --git a/modules/extensions.py b/modules/extensions.py
index db9c4200..b522125c 100644
--- a/modules/extensions.py
+++ b/modules/extensions.py
@@ -8,6 +8,7 @@ from modules import paths, shared
 

 extensions = []

 extensions_dir = os.path.join(paths.script_path, "extensions")

+extensions_builtin_dir = os.path.join(paths.script_path, "extensions-builtin")

 

 

 def active():

@@ -15,12 +16,13 @@ def active():
 

 

 class Extension:

-    def __init__(self, name, path, enabled=True):

+    def __init__(self, name, path, enabled=True, is_builtin=False):

         self.name = name

         self.path = path

         self.enabled = enabled

         self.status = ''

         self.can_update = False

+        self.is_builtin = is_builtin

 

         repo = None

         try:

@@ -79,11 +81,19 @@ def list_extensions():
     if not os.path.isdir(extensions_dir):

         return

 

-    for dirname in sorted(os.listdir(extensions_dir)):

-        path = os.path.join(extensions_dir, dirname)

-        if not os.path.isdir(path):

-            continue

+    paths = []

+    for dirname in [extensions_dir, extensions_builtin_dir]:

+        if not os.path.isdir(dirname):

+            return

 

-        extension = Extension(name=dirname, path=path, enabled=dirname not in shared.opts.disabled_extensions)

+        for extension_dirname in sorted(os.listdir(dirname)):

+            path = os.path.join(dirname, extension_dirname)

+            if not os.path.isdir(path):

+                continue

+

+            paths.append((extension_dirname, path, dirname == extensions_builtin_dir))

+

+    for dirname, path, is_builtin in paths:

+        extension = Extension(name=dirname, path=path, enabled=dirname not in shared.opts.disabled_extensions, is_builtin=is_builtin)

         extensions.append(extension)

 

diff --git a/modules/ldsr_model.py b/modules/ldsr_model.py
deleted file mode 100644
index 8c4db44a..00000000
--- a/modules/ldsr_model.py
+++ /dev/null
@@ -1,54 +0,0 @@
-import os
-import sys
-import traceback
-
-from basicsr.utils.download_util import load_file_from_url
-
-from modules.upscaler import Upscaler, UpscalerData
-from modules.ldsr_model_arch import LDSR
-from modules import shared
-
-
-class UpscalerLDSR(Upscaler):
-    def __init__(self, user_path):
-        self.name = "LDSR"
-        self.user_path = user_path
-        self.model_url = "https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1"
-        self.yaml_url = "https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1"
-        super().__init__()
-        scaler_data = UpscalerData("LDSR", None, self)
-        self.scalers = [scaler_data]
-
-    def load_model(self, path: str):
-        # Remove incorrect project.yaml file if too big
-        yaml_path = os.path.join(self.model_path, "project.yaml")
-        old_model_path = os.path.join(self.model_path, "model.pth")
-        new_model_path = os.path.join(self.model_path, "model.ckpt")
-        if os.path.exists(yaml_path):
-            statinfo = os.stat(yaml_path)
-            if statinfo.st_size >= 10485760:
-                print("Removing invalid LDSR YAML file.")
-                os.remove(yaml_path)
-        if os.path.exists(old_model_path):
-            print("Renaming model from model.pth to model.ckpt")
-            os.rename(old_model_path, new_model_path)
-        model = load_file_from_url(url=self.model_url, model_dir=self.model_path,
-                                   file_name="model.ckpt", progress=True)
-        yaml = load_file_from_url(url=self.yaml_url, model_dir=self.model_path,
-                                  file_name="project.yaml", progress=True)
-
-        try:
-            return LDSR(model, yaml)
-
-        except Exception:
-            print("Error importing LDSR:", file=sys.stderr)
-            print(traceback.format_exc(), file=sys.stderr)
-        return None
-
-    def do_upscale(self, img, path):
-        ldsr = self.load_model(path)
-        if ldsr is None:
-            print("NO LDSR!")
-            return img
-        ddim_steps = shared.opts.ldsr_steps
-        return ldsr.super_resolution(img, ddim_steps, self.scale)
diff --git a/modules/ldsr_model_arch.py b/modules/ldsr_model_arch.py
deleted file mode 100644
index 90e0a2f0..00000000
--- a/modules/ldsr_model_arch.py
+++ /dev/null
@@ -1,230 +0,0 @@
-import gc
-import time
-import warnings
-
-import numpy as np
-import torch
-import torchvision
-from PIL import Image
-from einops import rearrange, repeat
-from omegaconf import OmegaConf
-
-from ldm.models.diffusion.ddim import DDIMSampler
-from ldm.util import instantiate_from_config, ismap
-
-warnings.filterwarnings("ignore", category=UserWarning)
-
-
-# Create LDSR Class
-class LDSR:
-    def load_model_from_config(self, half_attention):
-        print(f"Loading model from {self.modelPath}")
-        pl_sd = torch.load(self.modelPath, map_location="cpu")
-        sd = pl_sd["state_dict"]
-        config = OmegaConf.load(self.yamlPath)
-        model = instantiate_from_config(config.model)
-        model.load_state_dict(sd, strict=False)
-        model.cuda()
-        if half_attention:
-            model = model.half()
-
-        model.eval()
-        return {"model": model}
-
-    def __init__(self, model_path, yaml_path):
-        self.modelPath = model_path
-        self.yamlPath = yaml_path
-
-    @staticmethod
-    def run(model, selected_path, custom_steps, eta):
-        example = get_cond(selected_path)
-
-        n_runs = 1
-        guider = None
-        ckwargs = None
-        ddim_use_x0_pred = False
-        temperature = 1.
-        eta = eta
-        custom_shape = None
-
-        height, width = example["image"].shape[1:3]
-        split_input = height >= 128 and width >= 128
-
-        if split_input:
-            ks = 128
-            stride = 64
-            vqf = 4  #
-            model.split_input_params = {"ks": (ks, ks), "stride": (stride, stride),
-                                        "vqf": vqf,
-                                        "patch_distributed_vq": True,
-                                        "tie_braker": False,
-                                        "clip_max_weight": 0.5,
-                                        "clip_min_weight": 0.01,
-                                        "clip_max_tie_weight": 0.5,
-                                        "clip_min_tie_weight": 0.01}
-        else:
-            if hasattr(model, "split_input_params"):
-                delattr(model, "split_input_params")
-
-        x_t = None
-        logs = None
-        for n in range(n_runs):
-            if custom_shape is not None:
-                x_t = torch.randn(1, custom_shape[1], custom_shape[2], custom_shape[3]).to(model.device)
-                x_t = repeat(x_t, '1 c h w -> b c h w', b=custom_shape[0])
-
-            logs = make_convolutional_sample(example, model,
-                                             custom_steps=custom_steps,
-                                             eta=eta, quantize_x0=False,
-                                             custom_shape=custom_shape,
-                                             temperature=temperature, noise_dropout=0.,
-                                             corrector=guider, corrector_kwargs=ckwargs, x_T=x_t,
-                                             ddim_use_x0_pred=ddim_use_x0_pred
-                                             )
-        return logs
-
-    def super_resolution(self, image, steps=100, target_scale=2, half_attention=False):
-        model = self.load_model_from_config(half_attention)
-
-        # Run settings
-        diffusion_steps = int(steps)
-        eta = 1.0
-
-        down_sample_method = 'Lanczos'
-
-        gc.collect()
-        torch.cuda.empty_cache()
-
-        im_og = image
-        width_og, height_og = im_og.size
-        # If we can adjust the max upscale size, then the 4 below should be our variable
-        down_sample_rate = target_scale / 4
-        wd = width_og * down_sample_rate
-        hd = height_og * down_sample_rate
-        width_downsampled_pre = int(np.ceil(wd))
-        height_downsampled_pre = int(np.ceil(hd))
-
-        if down_sample_rate != 1:
-            print(
-                f'Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]')
-            im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS)
-        else:
-            print(f"Down sample rate is 1 from {target_scale} / 4 (Not downsampling)")
-        
-        # pad width and height to multiples of 64, pads with the edge values of image to avoid artifacts
-        pad_w, pad_h = np.max(((2, 2), np.ceil(np.array(im_og.size) / 64).astype(int)), axis=0) * 64 - im_og.size
-        im_padded = Image.fromarray(np.pad(np.array(im_og), ((0, pad_h), (0, pad_w), (0, 0)), mode='edge'))
-        
-        logs = self.run(model["model"], im_padded, diffusion_steps, eta)
-
-        sample = logs["sample"]
-        sample = sample.detach().cpu()
-        sample = torch.clamp(sample, -1., 1.)
-        sample = (sample + 1.) / 2. * 255
-        sample = sample.numpy().astype(np.uint8)
-        sample = np.transpose(sample, (0, 2, 3, 1))
-        a = Image.fromarray(sample[0])
-
-        # remove padding
-        a = a.crop((0, 0) + tuple(np.array(im_og.size) * 4))
-
-        del model
-        gc.collect()
-        torch.cuda.empty_cache()
-        return a
-
-
-def get_cond(selected_path):
-    example = dict()
-    up_f = 4
-    c = selected_path.convert('RGB')
-    c = torch.unsqueeze(torchvision.transforms.ToTensor()(c), 0)
-    c_up = torchvision.transforms.functional.resize(c, size=[up_f * c.shape[2], up_f * c.shape[3]],
-                                                    antialias=True)
-    c_up = rearrange(c_up, '1 c h w -> 1 h w c')
-    c = rearrange(c, '1 c h w -> 1 h w c')
-    c = 2. * c - 1.
-
-    c = c.to(torch.device("cuda"))
-    example["LR_image"] = c
-    example["image"] = c_up
-
-    return example
-
-
-@torch.no_grad()
-def convsample_ddim(model, cond, steps, shape, eta=1.0, callback=None, normals_sequence=None,
-                    mask=None, x0=None, quantize_x0=False, temperature=1., score_corrector=None,
-                    corrector_kwargs=None, x_t=None
-                    ):
-    ddim = DDIMSampler(model)
-    bs = shape[0]
-    shape = shape[1:]
-    print(f"Sampling with eta = {eta}; steps: {steps}")
-    samples, intermediates = ddim.sample(steps, batch_size=bs, shape=shape, conditioning=cond, callback=callback,
-                                         normals_sequence=normals_sequence, quantize_x0=quantize_x0, eta=eta,
-                                         mask=mask, x0=x0, temperature=temperature, verbose=False,
-                                         score_corrector=score_corrector,
-                                         corrector_kwargs=corrector_kwargs, x_t=x_t)
-
-    return samples, intermediates
-
-
-@torch.no_grad()
-def make_convolutional_sample(batch, model, custom_steps=None, eta=1.0, quantize_x0=False, custom_shape=None, temperature=1., noise_dropout=0., corrector=None,
-                              corrector_kwargs=None, x_T=None, ddim_use_x0_pred=False):
-    log = dict()
-
-    z, c, x, xrec, xc = model.get_input(batch, model.first_stage_key,
-                                        return_first_stage_outputs=True,
-                                        force_c_encode=not (hasattr(model, 'split_input_params')
-                                                            and model.cond_stage_key == 'coordinates_bbox'),
-                                        return_original_cond=True)
-
-    if custom_shape is not None:
-        z = torch.randn(custom_shape)
-        print(f"Generating {custom_shape[0]} samples of shape {custom_shape[1:]}")
-
-    z0 = None
-
-    log["input"] = x
-    log["reconstruction"] = xrec
-
-    if ismap(xc):
-        log["original_conditioning"] = model.to_rgb(xc)
-        if hasattr(model, 'cond_stage_key'):
-            log[model.cond_stage_key] = model.to_rgb(xc)
-
-    else:
-        log["original_conditioning"] = xc if xc is not None else torch.zeros_like(x)
-        if model.cond_stage_model:
-            log[model.cond_stage_key] = xc if xc is not None else torch.zeros_like(x)
-            if model.cond_stage_key == 'class_label':
-                log[model.cond_stage_key] = xc[model.cond_stage_key]
-
-    with model.ema_scope("Plotting"):
-        t0 = time.time()
-
-        sample, intermediates = convsample_ddim(model, c, steps=custom_steps, shape=z.shape,
-                                                eta=eta,
-                                                quantize_x0=quantize_x0, mask=None, x0=z0,
-                                                temperature=temperature, score_corrector=corrector, corrector_kwargs=corrector_kwargs,
-                                                x_t=x_T)
-        t1 = time.time()
-
-        if ddim_use_x0_pred:
-            sample = intermediates['pred_x0'][-1]
-
-    x_sample = model.decode_first_stage(sample)
-
-    try:
-        x_sample_noquant = model.decode_first_stage(sample, force_not_quantize=True)
-        log["sample_noquant"] = x_sample_noquant
-        log["sample_diff"] = torch.abs(x_sample_noquant - x_sample)
-    except:
-        pass
-
-    log["sample"] = x_sample
-    log["time"] = t1 - t0
-
-    return log
diff --git a/modules/modelloader.py b/modules/modelloader.py
index 7d2f0ade..e647f6fa 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -124,10 +124,9 @@ def move_files(src_path: str, dest_path: str, ext_filter: str = None):
 
 
 def load_upscalers():
-    sd = shared.script_path
     # We can only do this 'magic' method to dynamically load upscalers if they are referenced,
     # so we'll try to import any _model.py files before looking in __subclasses__
-    modules_dir = os.path.join(sd, "modules")
+    modules_dir = os.path.join(shared.script_path, "modules")
     for file in os.listdir(modules_dir):
         if "_model.py" in file:
             model_name = file.replace("_model.py", "")
@@ -136,22 +135,13 @@ def load_upscalers():
                 importlib.import_module(full_model)
             except:
                 pass
+
     datas = []
-    c_o = vars(shared.cmd_opts)
+    commandline_options = vars(shared.cmd_opts)
     for cls in Upscaler.__subclasses__():
         name = cls.__name__
-        module_name = cls.__module__
-        module = importlib.import_module(module_name)
-        class_ = getattr(module, name)
         cmd_name = f"{name.lower().replace('upscaler', '')}_models_path"
-        opt_string = None
-        try:
-            if cmd_name in c_o:
-                opt_string = c_o[cmd_name]
-        except:
-            pass
-        scaler = class_(opt_string)
-        for child in scaler.scalers:
-            datas.append(child)
+        scaler = cls(commandline_options.get(cmd_name, None))
+        datas += scaler.scalers
 
     shared.sd_upscalers = datas
diff --git a/modules/scunet_model.py b/modules/scunet_model.py
deleted file mode 100644
index 52360241..00000000
--- a/modules/scunet_model.py
+++ /dev/null
@@ -1,87 +0,0 @@
-import os.path
-import sys
-import traceback
-
-import PIL.Image
-import numpy as np
-import torch
-from basicsr.utils.download_util import load_file_from_url
-
-import modules.upscaler
-from modules import devices, modelloader
-from modules.scunet_model_arch import SCUNet as net
-
-
-class UpscalerScuNET(modules.upscaler.Upscaler):
-    def __init__(self, dirname):
-        self.name = "ScuNET"
-        self.model_name = "ScuNET GAN"
-        self.model_name2 = "ScuNET PSNR"
-        self.model_url = "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_gan.pth"
-        self.model_url2 = "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth"
-        self.user_path = dirname
-        super().__init__()
-        model_paths = self.find_models(ext_filter=[".pth"])
-        scalers = []
-        add_model2 = True
-        for file in model_paths:
-            if "http" in file:
-                name = self.model_name
-            else:
-                name = modelloader.friendly_name(file)
-            if name == self.model_name2 or file == self.model_url2:
-                add_model2 = False
-            try:
-                scaler_data = modules.upscaler.UpscalerData(name, file, self, 4)
-                scalers.append(scaler_data)
-            except Exception:
-                print(f"Error loading ScuNET model: {file}", file=sys.stderr)
-                print(traceback.format_exc(), file=sys.stderr)
-        if add_model2:
-            scaler_data2 = modules.upscaler.UpscalerData(self.model_name2, self.model_url2, self)
-            scalers.append(scaler_data2)
-        self.scalers = scalers
-
-    def do_upscale(self, img: PIL.Image, selected_file):
-        torch.cuda.empty_cache()
-
-        model = self.load_model(selected_file)
-        if model is None:
-            return img
-
-        device = devices.device_scunet
-        img = np.array(img)
-        img = img[:, :, ::-1]
-        img = np.moveaxis(img, 2, 0) / 255
-        img = torch.from_numpy(img).float()
-        img = img.unsqueeze(0).to(device)
-
-        with torch.no_grad():
-            output = model(img)
-        output = output.squeeze().float().cpu().clamp_(0, 1).numpy()
-        output = 255. * np.moveaxis(output, 0, 2)
-        output = output.astype(np.uint8)
-        output = output[:, :, ::-1]
-        torch.cuda.empty_cache()
-        return PIL.Image.fromarray(output, 'RGB')
-
-    def load_model(self, path: str):
-        device = devices.device_scunet
-        if "http" in path:
-            filename = load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="%s.pth" % self.name,
-                                          progress=True)
-        else:
-            filename = path
-        if not os.path.exists(os.path.join(self.model_path, filename)) or filename is None:
-            print(f"ScuNET: Unable to load model from {filename}", file=sys.stderr)
-            return None
-
-        model = net(in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
-        model.load_state_dict(torch.load(filename), strict=True)
-        model.eval()
-        for k, v in model.named_parameters():
-            v.requires_grad = False
-        model = model.to(device)
-
-        return model
-
diff --git a/modules/scunet_model_arch.py b/modules/scunet_model_arch.py
deleted file mode 100644
index 43ca8d36..00000000
--- a/modules/scunet_model_arch.py
+++ /dev/null
@@ -1,265 +0,0 @@
-# -*- coding: utf-8 -*-
-import numpy as np
-import torch
-import torch.nn as nn
-from einops import rearrange
-from einops.layers.torch import Rearrange
-from timm.models.layers import trunc_normal_, DropPath
-
-
-class WMSA(nn.Module):
-    """ Self-attention module in Swin Transformer
-    """
-
-    def __init__(self, input_dim, output_dim, head_dim, window_size, type):
-        super(WMSA, self).__init__()
-        self.input_dim = input_dim
-        self.output_dim = output_dim
-        self.head_dim = head_dim
-        self.scale = self.head_dim ** -0.5
-        self.n_heads = input_dim // head_dim
-        self.window_size = window_size
-        self.type = type
-        self.embedding_layer = nn.Linear(self.input_dim, 3 * self.input_dim, bias=True)
-
-        self.relative_position_params = nn.Parameter(
-            torch.zeros((2 * window_size - 1) * (2 * window_size - 1), self.n_heads))
-
-        self.linear = nn.Linear(self.input_dim, self.output_dim)
-
-        trunc_normal_(self.relative_position_params, std=.02)
-        self.relative_position_params = torch.nn.Parameter(
-            self.relative_position_params.view(2 * window_size - 1, 2 * window_size - 1, self.n_heads).transpose(1,
-                                                                                                                 2).transpose(
-                0, 1))
-
-    def generate_mask(self, h, w, p, shift):
-        """ generating the mask of SW-MSA
-        Args:
-            shift: shift parameters in CyclicShift.
-        Returns:
-            attn_mask: should be (1 1 w p p),
-        """
-        # supporting square.
-        attn_mask = torch.zeros(h, w, p, p, p, p, dtype=torch.bool, device=self.relative_position_params.device)
-        if self.type == 'W':
-            return attn_mask
-
-        s = p - shift
-        attn_mask[-1, :, :s, :, s:, :] = True
-        attn_mask[-1, :, s:, :, :s, :] = True
-        attn_mask[:, -1, :, :s, :, s:] = True
-        attn_mask[:, -1, :, s:, :, :s] = True
-        attn_mask = rearrange(attn_mask, 'w1 w2 p1 p2 p3 p4 -> 1 1 (w1 w2) (p1 p2) (p3 p4)')
-        return attn_mask
-
-    def forward(self, x):
-        """ Forward pass of Window Multi-head Self-attention module.
-        Args:
-            x: input tensor with shape of [b h w c];
-            attn_mask: attention mask, fill -inf where the value is True;
-        Returns:
-            output: tensor shape [b h w c]
-        """
-        if self.type != 'W': x = torch.roll(x, shifts=(-(self.window_size // 2), -(self.window_size // 2)), dims=(1, 2))
-        x = rearrange(x, 'b (w1 p1) (w2 p2) c -> b w1 w2 p1 p2 c', p1=self.window_size, p2=self.window_size)
-        h_windows = x.size(1)
-        w_windows = x.size(2)
-        # square validation
-        # assert h_windows == w_windows
-
-        x = rearrange(x, 'b w1 w2 p1 p2 c -> b (w1 w2) (p1 p2) c', p1=self.window_size, p2=self.window_size)
-        qkv = self.embedding_layer(x)
-        q, k, v = rearrange(qkv, 'b nw np (threeh c) -> threeh b nw np c', c=self.head_dim).chunk(3, dim=0)
-        sim = torch.einsum('hbwpc,hbwqc->hbwpq', q, k) * self.scale
-        # Adding learnable relative embedding
-        sim = sim + rearrange(self.relative_embedding(), 'h p q -> h 1 1 p q')
-        # Using Attn Mask to distinguish different subwindows.
-        if self.type != 'W':
-            attn_mask = self.generate_mask(h_windows, w_windows, self.window_size, shift=self.window_size // 2)
-            sim = sim.masked_fill_(attn_mask, float("-inf"))
-
-        probs = nn.functional.softmax(sim, dim=-1)
-        output = torch.einsum('hbwij,hbwjc->hbwic', probs, v)
-        output = rearrange(output, 'h b w p c -> b w p (h c)')
-        output = self.linear(output)
-        output = rearrange(output, 'b (w1 w2) (p1 p2) c -> b (w1 p1) (w2 p2) c', w1=h_windows, p1=self.window_size)
-
-        if self.type != 'W': output = torch.roll(output, shifts=(self.window_size // 2, self.window_size // 2),
-                                                 dims=(1, 2))
-        return output
-
-    def relative_embedding(self):
-        cord = torch.tensor(np.array([[i, j] for i in range(self.window_size) for j in range(self.window_size)]))
-        relation = cord[:, None, :] - cord[None, :, :] + self.window_size - 1
-        # negative is allowed
-        return self.relative_position_params[:, relation[:, :, 0].long(), relation[:, :, 1].long()]
-
-
-class Block(nn.Module):
-    def __init__(self, input_dim, output_dim, head_dim, window_size, drop_path, type='W', input_resolution=None):
-        """ SwinTransformer Block
-        """
-        super(Block, self).__init__()
-        self.input_dim = input_dim
-        self.output_dim = output_dim
-        assert type in ['W', 'SW']
-        self.type = type
-        if input_resolution <= window_size:
-            self.type = 'W'
-
-        self.ln1 = nn.LayerNorm(input_dim)
-        self.msa = WMSA(input_dim, input_dim, head_dim, window_size, self.type)
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
-        self.ln2 = nn.LayerNorm(input_dim)
-        self.mlp = nn.Sequential(
-            nn.Linear(input_dim, 4 * input_dim),
-            nn.GELU(),
-            nn.Linear(4 * input_dim, output_dim),
-        )
-
-    def forward(self, x):
-        x = x + self.drop_path(self.msa(self.ln1(x)))
-        x = x + self.drop_path(self.mlp(self.ln2(x)))
-        return x
-
-
-class ConvTransBlock(nn.Module):
-    def __init__(self, conv_dim, trans_dim, head_dim, window_size, drop_path, type='W', input_resolution=None):
-        """ SwinTransformer and Conv Block
-        """
-        super(ConvTransBlock, self).__init__()
-        self.conv_dim = conv_dim
-        self.trans_dim = trans_dim
-        self.head_dim = head_dim
-        self.window_size = window_size
-        self.drop_path = drop_path
-        self.type = type
-        self.input_resolution = input_resolution
-
-        assert self.type in ['W', 'SW']
-        if self.input_resolution <= self.window_size:
-            self.type = 'W'
-
-        self.trans_block = Block(self.trans_dim, self.trans_dim, self.head_dim, self.window_size, self.drop_path,
-                                 self.type, self.input_resolution)
-        self.conv1_1 = nn.Conv2d(self.conv_dim + self.trans_dim, self.conv_dim + self.trans_dim, 1, 1, 0, bias=True)
-        self.conv1_2 = nn.Conv2d(self.conv_dim + self.trans_dim, self.conv_dim + self.trans_dim, 1, 1, 0, bias=True)
-
-        self.conv_block = nn.Sequential(
-            nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False),
-            nn.ReLU(True),
-            nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False)
-        )
-
-    def forward(self, x):
-        conv_x, trans_x = torch.split(self.conv1_1(x), (self.conv_dim, self.trans_dim), dim=1)
-        conv_x = self.conv_block(conv_x) + conv_x
-        trans_x = Rearrange('b c h w -> b h w c')(trans_x)
-        trans_x = self.trans_block(trans_x)
-        trans_x = Rearrange('b h w c -> b c h w')(trans_x)
-        res = self.conv1_2(torch.cat((conv_x, trans_x), dim=1))
-        x = x + res
-
-        return x
-
-
-class SCUNet(nn.Module):
-    # def __init__(self, in_nc=3, config=[2, 2, 2, 2, 2, 2, 2], dim=64, drop_path_rate=0.0, input_resolution=256):
-    def __init__(self, in_nc=3, config=None, dim=64, drop_path_rate=0.0, input_resolution=256):
-        super(SCUNet, self).__init__()
-        if config is None:
-            config = [2, 2, 2, 2, 2, 2, 2]
-        self.config = config
-        self.dim = dim
-        self.head_dim = 32
-        self.window_size = 8
-
-        # drop path rate for each layer
-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(config))]
-
-        self.m_head = [nn.Conv2d(in_nc, dim, 3, 1, 1, bias=False)]
-
-        begin = 0
-        self.m_down1 = [ConvTransBlock(dim // 2, dim // 2, self.head_dim, self.window_size, dpr[i + begin],
-                                       'W' if not i % 2 else 'SW', input_resolution)
-                        for i in range(config[0])] + \
-                       [nn.Conv2d(dim, 2 * dim, 2, 2, 0, bias=False)]
-
-        begin += config[0]
-        self.m_down2 = [ConvTransBlock(dim, dim, self.head_dim, self.window_size, dpr[i + begin],
-                                       'W' if not i % 2 else 'SW', input_resolution // 2)
-                        for i in range(config[1])] + \
-                       [nn.Conv2d(2 * dim, 4 * dim, 2, 2, 0, bias=False)]
-
-        begin += config[1]
-        self.m_down3 = [ConvTransBlock(2 * dim, 2 * dim, self.head_dim, self.window_size, dpr[i + begin],
-                                       'W' if not i % 2 else 'SW', input_resolution // 4)
-                        for i in range(config[2])] + \
-                       [nn.Conv2d(4 * dim, 8 * dim, 2, 2, 0, bias=False)]
-
-        begin += config[2]
-        self.m_body = [ConvTransBlock(4 * dim, 4 * dim, self.head_dim, self.window_size, dpr[i + begin],
-                                      'W' if not i % 2 else 'SW', input_resolution // 8)
-                       for i in range(config[3])]
-
-        begin += config[3]
-        self.m_up3 = [nn.ConvTranspose2d(8 * dim, 4 * dim, 2, 2, 0, bias=False), ] + \
-                     [ConvTransBlock(2 * dim, 2 * dim, self.head_dim, self.window_size, dpr[i + begin],
-                                     'W' if not i % 2 else 'SW', input_resolution // 4)
-                      for i in range(config[4])]
-
-        begin += config[4]
-        self.m_up2 = [nn.ConvTranspose2d(4 * dim, 2 * dim, 2, 2, 0, bias=False), ] + \
-                     [ConvTransBlock(dim, dim, self.head_dim, self.window_size, dpr[i + begin],
-                                     'W' if not i % 2 else 'SW', input_resolution // 2)
-                      for i in range(config[5])]
-
-        begin += config[5]
-        self.m_up1 = [nn.ConvTranspose2d(2 * dim, dim, 2, 2, 0, bias=False), ] + \
-                     [ConvTransBlock(dim // 2, dim // 2, self.head_dim, self.window_size, dpr[i + begin],
-                                     'W' if not i % 2 else 'SW', input_resolution)
-                      for i in range(config[6])]
-
-        self.m_tail = [nn.Conv2d(dim, in_nc, 3, 1, 1, bias=False)]
-
-        self.m_head = nn.Sequential(*self.m_head)
-        self.m_down1 = nn.Sequential(*self.m_down1)
-        self.m_down2 = nn.Sequential(*self.m_down2)
-        self.m_down3 = nn.Sequential(*self.m_down3)
-        self.m_body = nn.Sequential(*self.m_body)
-        self.m_up3 = nn.Sequential(*self.m_up3)
-        self.m_up2 = nn.Sequential(*self.m_up2)
-        self.m_up1 = nn.Sequential(*self.m_up1)
-        self.m_tail = nn.Sequential(*self.m_tail)
-        # self.apply(self._init_weights)
-
-    def forward(self, x0):
-
-        h, w = x0.size()[-2:]
-        paddingBottom = int(np.ceil(h / 64) * 64 - h)
-        paddingRight = int(np.ceil(w / 64) * 64 - w)
-        x0 = nn.ReplicationPad2d((0, paddingRight, 0, paddingBottom))(x0)
-
-        x1 = self.m_head(x0)
-        x2 = self.m_down1(x1)
-        x3 = self.m_down2(x2)
-        x4 = self.m_down3(x3)
-        x = self.m_body(x4)
-        x = self.m_up3(x + x4)
-        x = self.m_up2(x + x3)
-        x = self.m_up1(x + x2)
-        x = self.m_tail(x + x1)
-
-        x = x[..., :h, :w]
-
-        return x
-
-    def _init_weights(self, m):
-        if isinstance(m, nn.Linear):
-            trunc_normal_(m.weight, std=.02)
-            if m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.LayerNorm):
-            nn.init.constant_(m.bias, 0)
-            nn.init.constant_(m.weight, 1.0)
\ No newline at end of file
diff --git a/modules/shared.py b/modules/shared.py
index 8202d8e5..dc45fcaa 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -50,9 +50,6 @@ parser.add_argument("--gfpgan-models-path", type=str, help="Path to directory wi
 parser.add_argument("--esrgan-models-path", type=str, help="Path to directory with ESRGAN model file(s).", default=os.path.join(models_path, 'ESRGAN'))

 parser.add_argument("--bsrgan-models-path", type=str, help="Path to directory with BSRGAN model file(s).", default=os.path.join(models_path, 'BSRGAN'))

 parser.add_argument("--realesrgan-models-path", type=str, help="Path to directory with RealESRGAN model file(s).", default=os.path.join(models_path, 'RealESRGAN'))

-parser.add_argument("--scunet-models-path", type=str, help="Path to directory with ScuNET model file(s).", default=os.path.join(models_path, 'ScuNET'))

-parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR model file(s).", default=os.path.join(models_path, 'SwinIR'))

-parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR model file(s).", default=os.path.join(models_path, 'LDSR'))

 parser.add_argument("--clip-models-path", type=str, help="Path to directory with CLIP model file(s).", default=None)

 parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers")

 parser.add_argument("--force-enable-xformers", action='store_true', help="enable xformers for cross attention layers regardless of whether the checking code thinks you can run it; do not make bug reports if this fails to work")

@@ -61,7 +58,7 @@ parser.add_argument("--opt-split-attention", action='store_true', help="force-en
 parser.add_argument("--opt-split-attention-invokeai", action='store_true', help="force-enables InvokeAI's cross-attention layer optimization. By default, it's on when cuda is unavailable.")

 parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")

 parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")

-parser.add_argument("--use-cpu", nargs='+',choices=['all', 'sd', 'interrogate', 'gfpgan', 'swinir', 'esrgan', 'scunet', 'codeformer'], help="use CPU as torch device for specified modules", default=[], type=str.lower)

+parser.add_argument("--use-cpu", nargs='+', help="use CPU as torch device for specified modules", default=[], type=str.lower)

 parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")

 parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None)

 parser.add_argument("--show-negative-prompt", action='store_true', help="does not do anything", default=False)

@@ -95,6 +92,7 @@ parser.add_argument("--tls-certfile", type=str, help="Partially enables TLS, req
 parser.add_argument("--server-name", type=str, help="Sets hostname of server", default=None)

 

 script_loading.preload_extensions(extensions.extensions_dir, parser)

+script_loading.preload_extensions(extensions.extensions_builtin_dir, parser)

 

 cmd_opts = parser.parse_args()

 

@@ -112,8 +110,8 @@ restricted_opts = {
 

 cmd_opts.disable_extension_access = (cmd_opts.share or cmd_opts.listen or cmd_opts.server_name) and not cmd_opts.enable_insecure_extension_access

 

-devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_swinir, devices.device_esrgan, devices.device_scunet, devices.device_codeformer = \

-(devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'swinir', 'esrgan', 'scunet', 'codeformer'])

+devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_esrgan, devices.device_codeformer = \

+    (devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'esrgan', 'codeformer'])

 

 device = devices.device

 weight_load_location = None if cmd_opts.lowram else "cpu"

@@ -326,9 +324,6 @@ options_templates.update(options_section(('upscaling', "Upscaling"), {
     "ESRGAN_tile": OptionInfo(192, "Tile size for ESRGAN upscalers. 0 = no tiling.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}),

     "ESRGAN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for ESRGAN upscalers. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),

     "realesrgan_enabled_models": OptionInfo(["R-ESRGAN 4x+", "R-ESRGAN 4x+ Anime6B"], "Select which Real-ESRGAN models to show in the web UI. (Requires restart)", gr.CheckboxGroup, lambda: {"choices": realesrgan_models_names()}),

-    "SWIN_tile": OptionInfo(192, "Tile size for all SwinIR.", gr.Slider, {"minimum": 16, "maximum": 512, "step": 16}),

-    "SWIN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for SwinIR. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),

-    "ldsr_steps": OptionInfo(100, "LDSR processing steps. Lower = faster", gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}),

     "upscaler_for_img2img": OptionInfo(None, "Upscaler for img2img", gr.Dropdown, lambda: {"choices": [x.name for x in sd_upscalers]}),

     "use_scale_latent_for_hires_fix": OptionInfo(False, "Upscale latent space image when doing hires. fix"),

 }))

diff --git a/modules/swinir_model.py b/modules/swinir_model.py
deleted file mode 100644
index 483eabd4..00000000
--- a/modules/swinir_model.py
+++ /dev/null
@@ -1,157 +0,0 @@
-import contextlib
-import os
-
-import numpy as np
-import torch
-from PIL import Image
-from basicsr.utils.download_util import load_file_from_url
-from tqdm import tqdm
-
-from modules import modelloader, devices
-from modules.shared import cmd_opts, opts
-from modules.swinir_model_arch import SwinIR as net
-from modules.swinir_model_arch_v2 import Swin2SR as net2
-from modules.upscaler import Upscaler, UpscalerData
-
-
-class UpscalerSwinIR(Upscaler):
-    def __init__(self, dirname):
-        self.name = "SwinIR"
-        self.model_url = "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0" \
-                         "/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR" \
-                         "-L_x4_GAN.pth "
-        self.model_name = "SwinIR 4x"
-        self.user_path = dirname
-        super().__init__()
-        scalers = []
-        model_files = self.find_models(ext_filter=[".pt", ".pth"])
-        for model in model_files:
-            if "http" in model:
-                name = self.model_name
-            else:
-                name = modelloader.friendly_name(model)
-            model_data = UpscalerData(name, model, self)
-            scalers.append(model_data)
-        self.scalers = scalers
-
-    def do_upscale(self, img, model_file):
-        model = self.load_model(model_file)
-        if model is None:
-            return img
-        model = model.to(devices.device_swinir)
-        img = upscale(img, model)
-        try:
-            torch.cuda.empty_cache()
-        except:
-            pass
-        return img
-
-    def load_model(self, path, scale=4):
-        if "http" in path:
-            dl_name = "%s%s" % (self.model_name.replace(" ", "_"), ".pth")
-            filename = load_file_from_url(url=path, model_dir=self.model_path, file_name=dl_name, progress=True)
-        else:
-            filename = path
-        if filename is None or not os.path.exists(filename):
-            return None
-        if filename.endswith(".v2.pth"):
-            model = net2(
-            upscale=scale,
-            in_chans=3,
-            img_size=64,
-            window_size=8,
-            img_range=1.0,
-            depths=[6, 6, 6, 6, 6, 6],
-            embed_dim=180,
-            num_heads=[6, 6, 6, 6, 6, 6],
-            mlp_ratio=2,
-            upsampler="nearest+conv",
-            resi_connection="1conv",
-            )
-            params = None
-        else:
-            model = net(
-                upscale=scale,
-                in_chans=3,
-                img_size=64,
-                window_size=8,
-                img_range=1.0,
-                depths=[6, 6, 6, 6, 6, 6, 6, 6, 6],
-                embed_dim=240,
-                num_heads=[8, 8, 8, 8, 8, 8, 8, 8, 8],
-                mlp_ratio=2,
-                upsampler="nearest+conv",
-                resi_connection="3conv",
-            )
-            params = "params_ema"
-
-        pretrained_model = torch.load(filename)
-        if params is not None:
-            model.load_state_dict(pretrained_model[params], strict=True)
-        else:
-            model.load_state_dict(pretrained_model, strict=True)
-        if not cmd_opts.no_half:
-            model = model.half()
-        return model
-
-
-def upscale(
-        img,
-        model,
-        tile=opts.SWIN_tile,
-        tile_overlap=opts.SWIN_tile_overlap,
-        window_size=8,
-        scale=4,
-):
-    img = np.array(img)
-    img = img[:, :, ::-1]
-    img = np.moveaxis(img, 2, 0) / 255
-    img = torch.from_numpy(img).float()
-    img = img.unsqueeze(0).to(devices.device_swinir)
-    with torch.no_grad(), devices.autocast():
-        _, _, h_old, w_old = img.size()
-        h_pad = (h_old // window_size + 1) * window_size - h_old
-        w_pad = (w_old // window_size + 1) * window_size - w_old
-        img = torch.cat([img, torch.flip(img, [2])], 2)[:, :, : h_old + h_pad, :]
-        img = torch.cat([img, torch.flip(img, [3])], 3)[:, :, :, : w_old + w_pad]
-        output = inference(img, model, tile, tile_overlap, window_size, scale)
-        output = output[..., : h_old * scale, : w_old * scale]
-        output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
-        if output.ndim == 3:
-            output = np.transpose(
-                output[[2, 1, 0], :, :], (1, 2, 0)
-            )  # CHW-RGB to HCW-BGR
-        output = (output * 255.0).round().astype(np.uint8)  # float32 to uint8
-        return Image.fromarray(output, "RGB")
-
-
-def inference(img, model, tile, tile_overlap, window_size, scale):
-    # test the image tile by tile
-    b, c, h, w = img.size()
-    tile = min(tile, h, w)
-    assert tile % window_size == 0, "tile size should be a multiple of window_size"
-    sf = scale
-
-    stride = tile - tile_overlap
-    h_idx_list = list(range(0, h - tile, stride)) + [h - tile]
-    w_idx_list = list(range(0, w - tile, stride)) + [w - tile]
-    E = torch.zeros(b, c, h * sf, w * sf, dtype=torch.half, device=devices.device_swinir).type_as(img)
-    W = torch.zeros_like(E, dtype=torch.half, device=devices.device_swinir)
-
-    with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="SwinIR tiles") as pbar:
-        for h_idx in h_idx_list:
-            for w_idx in w_idx_list:
-                in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
-                out_patch = model(in_patch)
-                out_patch_mask = torch.ones_like(out_patch)
-
-                E[
-                ..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
-                ].add_(out_patch)
-                W[
-                ..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
-                ].add_(out_patch_mask)
-                pbar.update(1)
-    output = E.div_(W)
-
-    return output
diff --git a/modules/swinir_model_arch.py b/modules/swinir_model_arch.py
deleted file mode 100644
index 863f42db..00000000
--- a/modules/swinir_model_arch.py
+++ /dev/null
@@ -1,867 +0,0 @@
-# -----------------------------------------------------------------------------------
-# SwinIR: Image Restoration Using Swin Transformer, https://arxiv.org/abs/2108.10257
-# Originally Written by Ze Liu, Modified by Jingyun Liang.
-# -----------------------------------------------------------------------------------
-
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint as checkpoint
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
-
-
-class Mlp(nn.Module):
-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features)
-        self.act = act_layer()
-        self.fc2 = nn.Linear(hidden_features, out_features)
-        self.drop = nn.Dropout(drop)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.drop(x)
-        x = self.fc2(x)
-        x = self.drop(x)
-        return x
-
-
-def window_partition(x, window_size):
-    """
-    Args:
-        x: (B, H, W, C)
-        window_size (int): window size
-
-    Returns:
-        windows: (num_windows*B, window_size, window_size, C)
-    """
-    B, H, W, C = x.shape
-    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
-    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
-    return windows
-
-
-def window_reverse(windows, window_size, H, W):
-    """
-    Args:
-        windows: (num_windows*B, window_size, window_size, C)
-        window_size (int): Window size
-        H (int): Height of image
-        W (int): Width of image
-
-    Returns:
-        x: (B, H, W, C)
-    """
-    B = int(windows.shape[0] / (H * W / window_size / window_size))
-    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
-    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
-    return x
-
-
-class WindowAttention(nn.Module):
-    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
-    It supports both of shifted and non-shifted window.
-
-    Args:
-        dim (int): Number of input channels.
-        window_size (tuple[int]): The height and width of the window.
-        num_heads (int): Number of attention heads.
-        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
-        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
-        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
-    """
-
-    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
-
-        super().__init__()
-        self.dim = dim
-        self.window_size = window_size  # Wh, Ww
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        self.scale = qk_scale or head_dim ** -0.5
-
-        # define a parameter table of relative position bias
-        self.relative_position_bias_table = nn.Parameter(
-            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
-
-        # get pair-wise relative position index for each token inside the window
-        coords_h = torch.arange(self.window_size[0])
-        coords_w = torch.arange(self.window_size[1])
-        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
-        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
-        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
-        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
-        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
-        relative_coords[:, :, 1] += self.window_size[1] - 1
-        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
-        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
-        self.register_buffer("relative_position_index", relative_position_index)
-
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.proj = nn.Linear(dim, dim)
-
-        self.proj_drop = nn.Dropout(proj_drop)
-
-        trunc_normal_(self.relative_position_bias_table, std=.02)
-        self.softmax = nn.Softmax(dim=-1)
-
-    def forward(self, x, mask=None):
-        """
-        Args:
-            x: input features with shape of (num_windows*B, N, C)
-            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
-        """
-        B_, N, C = x.shape
-        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
-
-        q = q * self.scale
-        attn = (q @ k.transpose(-2, -1))
-
-        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
-            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
-        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
-        attn = attn + relative_position_bias.unsqueeze(0)
-
-        if mask is not None:
-            nW = mask.shape[0]
-            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
-            attn = attn.view(-1, self.num_heads, N, N)
-            attn = self.softmax(attn)
-        else:
-            attn = self.softmax(attn)
-
-        attn = self.attn_drop(attn)
-
-        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-    def extra_repr(self) -> str:
-        return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
-
-    def flops(self, N):
-        # calculate flops for 1 window with token length of N
-        flops = 0
-        # qkv = self.qkv(x)
-        flops += N * self.dim * 3 * self.dim
-        # attn = (q @ k.transpose(-2, -1))
-        flops += self.num_heads * N * (self.dim // self.num_heads) * N
-        #  x = (attn @ v)
-        flops += self.num_heads * N * N * (self.dim // self.num_heads)
-        # x = self.proj(x)
-        flops += N * self.dim * self.dim
-        return flops
-
-
-class SwinTransformerBlock(nn.Module):
-    r""" Swin Transformer Block.
-
-    Args:
-        dim (int): Number of input channels.
-        input_resolution (tuple[int]): Input resolution.
-        num_heads (int): Number of attention heads.
-        window_size (int): Window size.
-        shift_size (int): Shift size for SW-MSA.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
-        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
-        drop (float, optional): Dropout rate. Default: 0.0
-        attn_drop (float, optional): Attention dropout rate. Default: 0.0
-        drop_path (float, optional): Stochastic depth rate. Default: 0.0
-        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
-        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
-    """
-
-    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
-                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
-                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
-        super().__init__()
-        self.dim = dim
-        self.input_resolution = input_resolution
-        self.num_heads = num_heads
-        self.window_size = window_size
-        self.shift_size = shift_size
-        self.mlp_ratio = mlp_ratio
-        if min(self.input_resolution) <= self.window_size:
-            # if window size is larger than input resolution, we don't partition windows
-            self.shift_size = 0
-            self.window_size = min(self.input_resolution)
-        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
-
-        self.norm1 = norm_layer(dim)
-        self.attn = WindowAttention(
-            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
-            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
-
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
-        self.norm2 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
-
-        if self.shift_size > 0:
-            attn_mask = self.calculate_mask(self.input_resolution)
-        else:
-            attn_mask = None
-
-        self.register_buffer("attn_mask", attn_mask)
-
-    def calculate_mask(self, x_size):
-        # calculate attention mask for SW-MSA
-        H, W = x_size
-        img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
-        h_slices = (slice(0, -self.window_size),
-                    slice(-self.window_size, -self.shift_size),
-                    slice(-self.shift_size, None))
-        w_slices = (slice(0, -self.window_size),
-                    slice(-self.window_size, -self.shift_size),
-                    slice(-self.shift_size, None))
-        cnt = 0
-        for h in h_slices:
-            for w in w_slices:
-                img_mask[:, h, w, :] = cnt
-                cnt += 1
-
-        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
-        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
-        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
-        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
-
-        return attn_mask
-
-    def forward(self, x, x_size):
-        H, W = x_size
-        B, L, C = x.shape
-        # assert L == H * W, "input feature has wrong size"
-
-        shortcut = x
-        x = self.norm1(x)
-        x = x.view(B, H, W, C)
-
-        # cyclic shift
-        if self.shift_size > 0:
-            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
-        else:
-            shifted_x = x
-
-        # partition windows
-        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
-        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
-
-        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
-        if self.input_resolution == x_size:
-            attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
-        else:
-            attn_windows = self.attn(x_windows, mask=self.calculate_mask(x_size).to(x.device))
-
-        # merge windows
-        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
-        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
-
-        # reverse cyclic shift
-        if self.shift_size > 0:
-            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
-        else:
-            x = shifted_x
-        x = x.view(B, H * W, C)
-
-        # FFN
-        x = shortcut + self.drop_path(x)
-        x = x + self.drop_path(self.mlp(self.norm2(x)))
-
-        return x
-
-    def extra_repr(self) -> str:
-        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
-               f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
-
-    def flops(self):
-        flops = 0
-        H, W = self.input_resolution
-        # norm1
-        flops += self.dim * H * W
-        # W-MSA/SW-MSA
-        nW = H * W / self.window_size / self.window_size
-        flops += nW * self.attn.flops(self.window_size * self.window_size)
-        # mlp
-        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
-        # norm2
-        flops += self.dim * H * W
-        return flops
-
-
-class PatchMerging(nn.Module):
-    r""" Patch Merging Layer.
-
-    Args:
-        input_resolution (tuple[int]): Resolution of input feature.
-        dim (int): Number of input channels.
-        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
-    """
-
-    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
-        super().__init__()
-        self.input_resolution = input_resolution
-        self.dim = dim
-        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
-        self.norm = norm_layer(4 * dim)
-
-    def forward(self, x):
-        """
-        x: B, H*W, C
-        """
-        H, W = self.input_resolution
-        B, L, C = x.shape
-        assert L == H * W, "input feature has wrong size"
-        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
-
-        x = x.view(B, H, W, C)
-
-        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
-        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
-        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
-        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
-        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
-        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
-
-        x = self.norm(x)
-        x = self.reduction(x)
-
-        return x
-
-    def extra_repr(self) -> str:
-        return f"input_resolution={self.input_resolution}, dim={self.dim}"
-
-    def flops(self):
-        H, W = self.input_resolution
-        flops = H * W * self.dim
-        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
-        return flops
-
-
-class BasicLayer(nn.Module):
-    """ A basic Swin Transformer layer for one stage.
-
-    Args:
-        dim (int): Number of input channels.
-        input_resolution (tuple[int]): Input resolution.
-        depth (int): Number of blocks.
-        num_heads (int): Number of attention heads.
-        window_size (int): Local window size.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
-        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
-        drop (float, optional): Dropout rate. Default: 0.0
-        attn_drop (float, optional): Attention dropout rate. Default: 0.0
-        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
-        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
-        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
-    """
-
-    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
-                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
-                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False):
-
-        super().__init__()
-        self.dim = dim
-        self.input_resolution = input_resolution
-        self.depth = depth
-        self.use_checkpoint = use_checkpoint
-
-        # build blocks
-        self.blocks = nn.ModuleList([
-            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
-                                 num_heads=num_heads, window_size=window_size,
-                                 shift_size=0 if (i % 2 == 0) else window_size // 2,
-                                 mlp_ratio=mlp_ratio,
-                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
-                                 drop=drop, attn_drop=attn_drop,
-                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
-                                 norm_layer=norm_layer)
-            for i in range(depth)])
-
-        # patch merging layer
-        if downsample is not None:
-            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
-        else:
-            self.downsample = None
-
-    def forward(self, x, x_size):
-        for blk in self.blocks:
-            if self.use_checkpoint:
-                x = checkpoint.checkpoint(blk, x, x_size)
-            else:
-                x = blk(x, x_size)
-        if self.downsample is not None:
-            x = self.downsample(x)
-        return x
-
-    def extra_repr(self) -> str:
-        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
-
-    def flops(self):
-        flops = 0
-        for blk in self.blocks:
-            flops += blk.flops()
-        if self.downsample is not None:
-            flops += self.downsample.flops()
-        return flops
-
-
-class RSTB(nn.Module):
-    """Residual Swin Transformer Block (RSTB).
-
-    Args:
-        dim (int): Number of input channels.
-        input_resolution (tuple[int]): Input resolution.
-        depth (int): Number of blocks.
-        num_heads (int): Number of attention heads.
-        window_size (int): Local window size.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
-        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
-        drop (float, optional): Dropout rate. Default: 0.0
-        attn_drop (float, optional): Attention dropout rate. Default: 0.0
-        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
-        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
-        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
-        img_size: Input image size.
-        patch_size: Patch size.
-        resi_connection: The convolutional block before residual connection.
-    """
-
-    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
-                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
-                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
-                 img_size=224, patch_size=4, resi_connection='1conv'):
-        super(RSTB, self).__init__()
-
-        self.dim = dim
-        self.input_resolution = input_resolution
-
-        self.residual_group = BasicLayer(dim=dim,
-                                         input_resolution=input_resolution,
-                                         depth=depth,
-                                         num_heads=num_heads,
-                                         window_size=window_size,
-                                         mlp_ratio=mlp_ratio,
-                                         qkv_bias=qkv_bias, qk_scale=qk_scale,
-                                         drop=drop, attn_drop=attn_drop,
-                                         drop_path=drop_path,
-                                         norm_layer=norm_layer,
-                                         downsample=downsample,
-                                         use_checkpoint=use_checkpoint)
-
-        if resi_connection == '1conv':
-            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
-        elif resi_connection == '3conv':
-            # to save parameters and memory
-            self.conv = nn.Sequential(nn.Conv2d(dim, dim // 4, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2, inplace=True),
-                                      nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),
-                                      nn.LeakyReLU(negative_slope=0.2, inplace=True),
-                                      nn.Conv2d(dim // 4, dim, 3, 1, 1))
-
-        self.patch_embed = PatchEmbed(
-            img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim,
-            norm_layer=None)
-
-        self.patch_unembed = PatchUnEmbed(
-            img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim,
-            norm_layer=None)
-
-    def forward(self, x, x_size):
-        return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))) + x
-
-    def flops(self):
-        flops = 0
-        flops += self.residual_group.flops()
-        H, W = self.input_resolution
-        flops += H * W * self.dim * self.dim * 9
-        flops += self.patch_embed.flops()
-        flops += self.patch_unembed.flops()
-
-        return flops
-
-
-class PatchEmbed(nn.Module):
-    r""" Image to Patch Embedding
-
-    Args:
-        img_size (int): Image size.  Default: 224.
-        patch_size (int): Patch token size. Default: 4.
-        in_chans (int): Number of input image channels. Default: 3.
-        embed_dim (int): Number of linear projection output channels. Default: 96.
-        norm_layer (nn.Module, optional): Normalization layer. Default: None
-    """
-
-    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
-        super().__init__()
-        img_size = to_2tuple(img_size)
-        patch_size = to_2tuple(patch_size)
-        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.patches_resolution = patches_resolution
-        self.num_patches = patches_resolution[0] * patches_resolution[1]
-
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
-
-        if norm_layer is not None:
-            self.norm = norm_layer(embed_dim)
-        else:
-            self.norm = None
-
-    def forward(self, x):
-        x = x.flatten(2).transpose(1, 2)  # B Ph*Pw C
-        if self.norm is not None:
-            x = self.norm(x)
-        return x
-
-    def flops(self):
-        flops = 0
-        H, W = self.img_size
-        if self.norm is not None:
-            flops += H * W * self.embed_dim
-        return flops
-
-
-class PatchUnEmbed(nn.Module):
-    r""" Image to Patch Unembedding
-
-    Args:
-        img_size (int): Image size.  Default: 224.
-        patch_size (int): Patch token size. Default: 4.
-        in_chans (int): Number of input image channels. Default: 3.
-        embed_dim (int): Number of linear projection output channels. Default: 96.
-        norm_layer (nn.Module, optional): Normalization layer. Default: None
-    """
-
-    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
-        super().__init__()
-        img_size = to_2tuple(img_size)
-        patch_size = to_2tuple(patch_size)
-        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.patches_resolution = patches_resolution
-        self.num_patches = patches_resolution[0] * patches_resolution[1]
-
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
-
-    def forward(self, x, x_size):
-        B, HW, C = x.shape
-        x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C
-        return x
-
-    def flops(self):
-        flops = 0
-        return flops
-
-
-class Upsample(nn.Sequential):
-    """Upsample module.
-
-    Args:
-        scale (int): Scale factor. Supported scales: 2^n and 3.
-        num_feat (int): Channel number of intermediate features.
-    """
-
-    def __init__(self, scale, num_feat):
-        m = []
-        if (scale & (scale - 1)) == 0:  # scale = 2^n
-            for _ in range(int(math.log(scale, 2))):
-                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
-                m.append(nn.PixelShuffle(2))
-        elif scale == 3:
-            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
-            m.append(nn.PixelShuffle(3))
-        else:
-            raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
-        super(Upsample, self).__init__(*m)
-
-
-class UpsampleOneStep(nn.Sequential):
-    """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
-       Used in lightweight SR to save parameters.
-
-    Args:
-        scale (int): Scale factor. Supported scales: 2^n and 3.
-        num_feat (int): Channel number of intermediate features.
-
-    """
-
-    def __init__(self, scale, num_feat, num_out_ch, input_resolution=None):
-        self.num_feat = num_feat
-        self.input_resolution = input_resolution
-        m = []
-        m.append(nn.Conv2d(num_feat, (scale ** 2) * num_out_ch, 3, 1, 1))
-        m.append(nn.PixelShuffle(scale))
-        super(UpsampleOneStep, self).__init__(*m)
-
-    def flops(self):
-        H, W = self.input_resolution
-        flops = H * W * self.num_feat * 3 * 9
-        return flops
-
-
-class SwinIR(nn.Module):
-    r""" SwinIR
-        A PyTorch impl of : `SwinIR: Image Restoration Using Swin Transformer`, based on Swin Transformer.
-
-    Args:
-        img_size (int | tuple(int)): Input image size. Default 64
-        patch_size (int | tuple(int)): Patch size. Default: 1
-        in_chans (int): Number of input image channels. Default: 3
-        embed_dim (int): Patch embedding dimension. Default: 96
-        depths (tuple(int)): Depth of each Swin Transformer layer.
-        num_heads (tuple(int)): Number of attention heads in different layers.
-        window_size (int): Window size. Default: 7
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
-        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
-        drop_rate (float): Dropout rate. Default: 0
-        attn_drop_rate (float): Attention dropout rate. Default: 0
-        drop_path_rate (float): Stochastic depth rate. Default: 0.1
-        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
-        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
-        patch_norm (bool): If True, add normalization after patch embedding. Default: True
-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
-        upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
-        img_range: Image range. 1. or 255.
-        upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
-        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
-    """
-
-    def __init__(self, img_size=64, patch_size=1, in_chans=3,
-                 embed_dim=96, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6],
-                 window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
-                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
-                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
-                 use_checkpoint=False, upscale=2, img_range=1., upsampler='', resi_connection='1conv',
-                 **kwargs):
-        super(SwinIR, self).__init__()
-        num_in_ch = in_chans
-        num_out_ch = in_chans
-        num_feat = 64
-        self.img_range = img_range
-        if in_chans == 3:
-            rgb_mean = (0.4488, 0.4371, 0.4040)
-            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
-        else:
-            self.mean = torch.zeros(1, 1, 1, 1)
-        self.upscale = upscale
-        self.upsampler = upsampler
-        self.window_size = window_size
-
-        #####################################################################################################
-        ################################### 1, shallow feature extraction ###################################
-        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)
-
-        #####################################################################################################
-        ################################### 2, deep feature extraction ######################################
-        self.num_layers = len(depths)
-        self.embed_dim = embed_dim
-        self.ape = ape
-        self.patch_norm = patch_norm
-        self.num_features = embed_dim
-        self.mlp_ratio = mlp_ratio
-
-        # split image into non-overlapping patches
-        self.patch_embed = PatchEmbed(
-            img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,
-            norm_layer=norm_layer if self.patch_norm else None)
-        num_patches = self.patch_embed.num_patches
-        patches_resolution = self.patch_embed.patches_resolution
-        self.patches_resolution = patches_resolution
-
-        # merge non-overlapping patches into image
-        self.patch_unembed = PatchUnEmbed(
-            img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,
-            norm_layer=norm_layer if self.patch_norm else None)
-
-        # absolute position embedding
-        if self.ape:
-            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
-            trunc_normal_(self.absolute_pos_embed, std=.02)
-
-        self.pos_drop = nn.Dropout(p=drop_rate)
-
-        # stochastic depth
-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
-
-        # build Residual Swin Transformer blocks (RSTB)
-        self.layers = nn.ModuleList()
-        for i_layer in range(self.num_layers):
-            layer = RSTB(dim=embed_dim,
-                         input_resolution=(patches_resolution[0],
-                                           patches_resolution[1]),
-                         depth=depths[i_layer],
-                         num_heads=num_heads[i_layer],
-                         window_size=window_size,
-                         mlp_ratio=self.mlp_ratio,
-                         qkv_bias=qkv_bias, qk_scale=qk_scale,
-                         drop=drop_rate, attn_drop=attn_drop_rate,
-                         drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],  # no impact on SR results
-                         norm_layer=norm_layer,
-                         downsample=None,
-                         use_checkpoint=use_checkpoint,
-                         img_size=img_size,
-                         patch_size=patch_size,
-                         resi_connection=resi_connection
-
-                         )
-            self.layers.append(layer)
-        self.norm = norm_layer(self.num_features)
-
-        # build the last conv layer in deep feature extraction
-        if resi_connection == '1conv':
-            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
-        elif resi_connection == '3conv':
-            # to save parameters and memory
-            self.conv_after_body = nn.Sequential(nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),
-                                                 nn.LeakyReLU(negative_slope=0.2, inplace=True),
-                                                 nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),
-                                                 nn.LeakyReLU(negative_slope=0.2, inplace=True),
-                                                 nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1))
-
-        #####################################################################################################
-        ################################ 3, high quality image reconstruction ################################
-        if self.upsampler == 'pixelshuffle':
-            # for classical SR
-            self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),
-                                                      nn.LeakyReLU(inplace=True))
-            self.upsample = Upsample(upscale, num_feat)
-            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
-        elif self.upsampler == 'pixelshuffledirect':
-            # for lightweight SR (to save parameters)
-            self.upsample = UpsampleOneStep(upscale, embed_dim, num_out_ch,
-                                            (patches_resolution[0], patches_resolution[1]))
-        elif self.upsampler == 'nearest+conv':
-            # for real-world SR (less artifacts)
-            self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),
-                                                      nn.LeakyReLU(inplace=True))
-            self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
-            if self.upscale == 4:
-                self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
-            self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
-            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
-            self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
-        else:
-            # for image denoising and JPEG compression artifact reduction
-            self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1)
-
-        self.apply(self._init_weights)
-
-    def _init_weights(self, m):
-        if isinstance(m, nn.Linear):
-            trunc_normal_(m.weight, std=.02)
-            if isinstance(m, nn.Linear) and m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.LayerNorm):
-            nn.init.constant_(m.bias, 0)
-            nn.init.constant_(m.weight, 1.0)
-
-    @torch.jit.ignore
-    def no_weight_decay(self):
-        return {'absolute_pos_embed'}
-
-    @torch.jit.ignore
-    def no_weight_decay_keywords(self):
-        return {'relative_position_bias_table'}
-
-    def check_image_size(self, x):
-        _, _, h, w = x.size()
-        mod_pad_h = (self.window_size - h % self.window_size) % self.window_size
-        mod_pad_w = (self.window_size - w % self.window_size) % self.window_size
-        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect')
-        return x
-
-    def forward_features(self, x):
-        x_size = (x.shape[2], x.shape[3])
-        x = self.patch_embed(x)
-        if self.ape:
-            x = x + self.absolute_pos_embed
-        x = self.pos_drop(x)
-
-        for layer in self.layers:
-            x = layer(x, x_size)
-
-        x = self.norm(x)  # B L C
-        x = self.patch_unembed(x, x_size)
-
-        return x
-
-    def forward(self, x):
-        H, W = x.shape[2:]
-        x = self.check_image_size(x)
-        
-        self.mean = self.mean.type_as(x)
-        x = (x - self.mean) * self.img_range
-
-        if self.upsampler == 'pixelshuffle':
-            # for classical SR
-            x = self.conv_first(x)
-            x = self.conv_after_body(self.forward_features(x)) + x
-            x = self.conv_before_upsample(x)
-            x = self.conv_last(self.upsample(x))
-        elif self.upsampler == 'pixelshuffledirect':
-            # for lightweight SR
-            x = self.conv_first(x)
-            x = self.conv_after_body(self.forward_features(x)) + x
-            x = self.upsample(x)
-        elif self.upsampler == 'nearest+conv':
-            # for real-world SR
-            x = self.conv_first(x)
-            x = self.conv_after_body(self.forward_features(x)) + x
-            x = self.conv_before_upsample(x)
-            x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
-            if self.upscale == 4:
-                x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
-            x = self.conv_last(self.lrelu(self.conv_hr(x)))
-        else:
-            # for image denoising and JPEG compression artifact reduction
-            x_first = self.conv_first(x)
-            res = self.conv_after_body(self.forward_features(x_first)) + x_first
-            x = x + self.conv_last(res)
-
-        x = x / self.img_range + self.mean
-
-        return x[:, :, :H*self.upscale, :W*self.upscale]
-
-    def flops(self):
-        flops = 0
-        H, W = self.patches_resolution
-        flops += H * W * 3 * self.embed_dim * 9
-        flops += self.patch_embed.flops()
-        for i, layer in enumerate(self.layers):
-            flops += layer.flops()
-        flops += H * W * 3 * self.embed_dim * self.embed_dim
-        flops += self.upsample.flops()
-        return flops
-
-
-if __name__ == '__main__':
-    upscale = 4
-    window_size = 8
-    height = (1024 // upscale // window_size + 1) * window_size
-    width = (720 // upscale // window_size + 1) * window_size
-    model = SwinIR(upscale=2, img_size=(height, width),
-                   window_size=window_size, img_range=1., depths=[6, 6, 6, 6],
-                   embed_dim=60, num_heads=[6, 6, 6, 6], mlp_ratio=2, upsampler='pixelshuffledirect')
-    print(model)
-    print(height, width, model.flops() / 1e9)
-
-    x = torch.randn((1, 3, height, width))
-    x = model(x)
-    print(x.shape)
diff --git a/modules/swinir_model_arch_v2.py b/modules/swinir_model_arch_v2.py
deleted file mode 100644
index 0e28ae6e..00000000
--- a/modules/swinir_model_arch_v2.py
+++ /dev/null
@@ -1,1017 +0,0 @@
-# -----------------------------------------------------------------------------------

-# Swin2SR: Swin2SR: SwinV2 Transformer for Compressed Image Super-Resolution and Restoration, https://arxiv.org/abs/

-# Written by Conde and Choi et al.

-# -----------------------------------------------------------------------------------

-

-import math

-import numpy as np

-import torch

-import torch.nn as nn

-import torch.nn.functional as F

-import torch.utils.checkpoint as checkpoint

-from timm.models.layers import DropPath, to_2tuple, trunc_normal_

-

-

-class Mlp(nn.Module):

-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):

-        super().__init__()

-        out_features = out_features or in_features

-        hidden_features = hidden_features or in_features

-        self.fc1 = nn.Linear(in_features, hidden_features)

-        self.act = act_layer()

-        self.fc2 = nn.Linear(hidden_features, out_features)

-        self.drop = nn.Dropout(drop)

-

-    def forward(self, x):

-        x = self.fc1(x)

-        x = self.act(x)

-        x = self.drop(x)

-        x = self.fc2(x)

-        x = self.drop(x)

-        return x

-

-

-def window_partition(x, window_size):

-    """

-    Args:

-        x: (B, H, W, C)

-        window_size (int): window size

-    Returns:

-        windows: (num_windows*B, window_size, window_size, C)

-    """

-    B, H, W, C = x.shape

-    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)

-    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)

-    return windows

-

-

-def window_reverse(windows, window_size, H, W):

-    """

-    Args:

-        windows: (num_windows*B, window_size, window_size, C)

-        window_size (int): Window size

-        H (int): Height of image

-        W (int): Width of image

-    Returns:

-        x: (B, H, W, C)

-    """

-    B = int(windows.shape[0] / (H * W / window_size / window_size))

-    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)

-    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)

-    return x

-

-class WindowAttention(nn.Module):

-    r""" Window based multi-head self attention (W-MSA) module with relative position bias.

-    It supports both of shifted and non-shifted window.

-    Args:

-        dim (int): Number of input channels.

-        window_size (tuple[int]): The height and width of the window.

-        num_heads (int): Number of attention heads.

-        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True

-        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0

-        proj_drop (float, optional): Dropout ratio of output. Default: 0.0

-        pretrained_window_size (tuple[int]): The height and width of the window in pre-training.

-    """

-

-    def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.,

-                 pretrained_window_size=[0, 0]):

-

-        super().__init__()

-        self.dim = dim

-        self.window_size = window_size  # Wh, Ww

-        self.pretrained_window_size = pretrained_window_size

-        self.num_heads = num_heads

-

-        self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)

-

-        # mlp to generate continuous relative position bias

-        self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True),

-                                     nn.ReLU(inplace=True),

-                                     nn.Linear(512, num_heads, bias=False))

-

-        # get relative_coords_table

-        relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32)

-        relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32)

-        relative_coords_table = torch.stack(

-            torch.meshgrid([relative_coords_h,

-                            relative_coords_w])).permute(1, 2, 0).contiguous().unsqueeze(0)  # 1, 2*Wh-1, 2*Ww-1, 2

-        if pretrained_window_size[0] > 0:

-            relative_coords_table[:, :, :, 0] /= (pretrained_window_size[0] - 1)

-            relative_coords_table[:, :, :, 1] /= (pretrained_window_size[1] - 1)

-        else:

-            relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1)

-            relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1)

-        relative_coords_table *= 8  # normalize to -8, 8

-        relative_coords_table = torch.sign(relative_coords_table) * torch.log2(

-            torch.abs(relative_coords_table) + 1.0) / np.log2(8)

-

-        self.register_buffer("relative_coords_table", relative_coords_table)

-

-        # get pair-wise relative position index for each token inside the window

-        coords_h = torch.arange(self.window_size[0])

-        coords_w = torch.arange(self.window_size[1])

-        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww

-        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww

-        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww

-        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2

-        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0

-        relative_coords[:, :, 1] += self.window_size[1] - 1

-        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1

-        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww

-        self.register_buffer("relative_position_index", relative_position_index)

-

-        self.qkv = nn.Linear(dim, dim * 3, bias=False)

-        if qkv_bias:

-            self.q_bias = nn.Parameter(torch.zeros(dim))

-            self.v_bias = nn.Parameter(torch.zeros(dim))

-        else:

-            self.q_bias = None

-            self.v_bias = None

-        self.attn_drop = nn.Dropout(attn_drop)

-        self.proj = nn.Linear(dim, dim)

-        self.proj_drop = nn.Dropout(proj_drop)

-        self.softmax = nn.Softmax(dim=-1)

-

-    def forward(self, x, mask=None):

-        """

-        Args:

-            x: input features with shape of (num_windows*B, N, C)

-            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None

-        """

-        B_, N, C = x.shape

-        qkv_bias = None

-        if self.q_bias is not None:

-            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))

-        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)

-        qkv = qkv.reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)

-        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)

-

-        # cosine attention

-        attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))

-        logit_scale = torch.clamp(self.logit_scale, max=torch.log(torch.tensor(1. / 0.01)).to(self.logit_scale.device)).exp()

-        attn = attn * logit_scale

-

-        relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads)

-        relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(

-            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH

-        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww

-        relative_position_bias = 16 * torch.sigmoid(relative_position_bias)

-        attn = attn + relative_position_bias.unsqueeze(0)

-

-        if mask is not None:

-            nW = mask.shape[0]

-            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)

-            attn = attn.view(-1, self.num_heads, N, N)

-            attn = self.softmax(attn)

-        else:

-            attn = self.softmax(attn)

-

-        attn = self.attn_drop(attn)

-

-        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)

-        x = self.proj(x)

-        x = self.proj_drop(x)

-        return x

-

-    def extra_repr(self) -> str:

-        return f'dim={self.dim}, window_size={self.window_size}, ' \

-               f'pretrained_window_size={self.pretrained_window_size}, num_heads={self.num_heads}'

-

-    def flops(self, N):

-        # calculate flops for 1 window with token length of N

-        flops = 0

-        # qkv = self.qkv(x)

-        flops += N * self.dim * 3 * self.dim

-        # attn = (q @ k.transpose(-2, -1))

-        flops += self.num_heads * N * (self.dim // self.num_heads) * N

-        #  x = (attn @ v)

-        flops += self.num_heads * N * N * (self.dim // self.num_heads)

-        # x = self.proj(x)

-        flops += N * self.dim * self.dim

-        return flops

-

-class SwinTransformerBlock(nn.Module):

-    r""" Swin Transformer Block.

-    Args:

-        dim (int): Number of input channels.

-        input_resolution (tuple[int]): Input resulotion.

-        num_heads (int): Number of attention heads.

-        window_size (int): Window size.

-        shift_size (int): Shift size for SW-MSA.

-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.

-        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True

-        drop (float, optional): Dropout rate. Default: 0.0

-        attn_drop (float, optional): Attention dropout rate. Default: 0.0

-        drop_path (float, optional): Stochastic depth rate. Default: 0.0

-        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU

-        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm

-        pretrained_window_size (int): Window size in pre-training.

-    """

-

-    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,

-                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., drop_path=0.,

-                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, pretrained_window_size=0):

-        super().__init__()

-        self.dim = dim

-        self.input_resolution = input_resolution

-        self.num_heads = num_heads

-        self.window_size = window_size

-        self.shift_size = shift_size

-        self.mlp_ratio = mlp_ratio

-        if min(self.input_resolution) <= self.window_size:

-            # if window size is larger than input resolution, we don't partition windows

-            self.shift_size = 0

-            self.window_size = min(self.input_resolution)

-        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"

-

-        self.norm1 = norm_layer(dim)

-        self.attn = WindowAttention(

-            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,

-            qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop,

-            pretrained_window_size=to_2tuple(pretrained_window_size))

-

-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

-        self.norm2 = norm_layer(dim)

-        mlp_hidden_dim = int(dim * mlp_ratio)

-        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

-

-        if self.shift_size > 0:

-            attn_mask = self.calculate_mask(self.input_resolution)

-        else:

-            attn_mask = None

-

-        self.register_buffer("attn_mask", attn_mask)

-        

-    def calculate_mask(self, x_size):

-        # calculate attention mask for SW-MSA

-        H, W = x_size

-        img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1

-        h_slices = (slice(0, -self.window_size),

-                    slice(-self.window_size, -self.shift_size),

-                    slice(-self.shift_size, None))

-        w_slices = (slice(0, -self.window_size),

-                    slice(-self.window_size, -self.shift_size),

-                    slice(-self.shift_size, None))

-        cnt = 0

-        for h in h_slices:

-            for w in w_slices:

-                img_mask[:, h, w, :] = cnt

-                cnt += 1

-

-        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1

-        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)

-        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)

-        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))

-

-        return attn_mask        

-

-    def forward(self, x, x_size):

-        H, W = x_size

-        B, L, C = x.shape

-        #assert L == H * W, "input feature has wrong size"

-

-        shortcut = x

-        x = x.view(B, H, W, C)

-

-        # cyclic shift

-        if self.shift_size > 0:

-            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))

-        else:

-            shifted_x = x

-

-        # partition windows

-        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C

-        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C

-

-        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size

-        if self.input_resolution == x_size:

-            attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C

-        else:

-            attn_windows = self.attn(x_windows, mask=self.calculate_mask(x_size).to(x.device))

-            

-        # merge windows

-        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)

-        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C

-

-        # reverse cyclic shift

-        if self.shift_size > 0:

-            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))

-        else:

-            x = shifted_x

-        x = x.view(B, H * W, C)

-        x = shortcut + self.drop_path(self.norm1(x))

-

-        # FFN

-        x = x + self.drop_path(self.norm2(self.mlp(x)))

-

-        return x

-

-    def extra_repr(self) -> str:

-        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \

-               f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"

-

-    def flops(self):

-        flops = 0

-        H, W = self.input_resolution

-        # norm1

-        flops += self.dim * H * W

-        # W-MSA/SW-MSA

-        nW = H * W / self.window_size / self.window_size

-        flops += nW * self.attn.flops(self.window_size * self.window_size)

-        # mlp

-        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio

-        # norm2

-        flops += self.dim * H * W

-        return flops

-

-class PatchMerging(nn.Module):

-    r""" Patch Merging Layer.

-    Args:

-        input_resolution (tuple[int]): Resolution of input feature.

-        dim (int): Number of input channels.

-        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm

-    """

-

-    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):

-        super().__init__()

-        self.input_resolution = input_resolution

-        self.dim = dim

-        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)

-        self.norm = norm_layer(2 * dim)

-

-    def forward(self, x):

-        """

-        x: B, H*W, C

-        """

-        H, W = self.input_resolution

-        B, L, C = x.shape

-        assert L == H * W, "input feature has wrong size"

-        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."

-

-        x = x.view(B, H, W, C)

-

-        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C

-        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C

-        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C

-        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C

-        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C

-        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C

-

-        x = self.reduction(x)

-        x = self.norm(x)

-

-        return x

-

-    def extra_repr(self) -> str:

-        return f"input_resolution={self.input_resolution}, dim={self.dim}"

-

-    def flops(self):

-        H, W = self.input_resolution

-        flops = (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim

-        flops += H * W * self.dim // 2

-        return flops    

-

-class BasicLayer(nn.Module):

-    """ A basic Swin Transformer layer for one stage.

-    Args:

-        dim (int): Number of input channels.

-        input_resolution (tuple[int]): Input resolution.

-        depth (int): Number of blocks.

-        num_heads (int): Number of attention heads.

-        window_size (int): Local window size.

-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.

-        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True

-        drop (float, optional): Dropout rate. Default: 0.0

-        attn_drop (float, optional): Attention dropout rate. Default: 0.0

-        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0

-        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm

-        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None

-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.

-        pretrained_window_size (int): Local window size in pre-training.

-    """

-

-    def __init__(self, dim, input_resolution, depth, num_heads, window_size,

-                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0.,

-                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,

-                 pretrained_window_size=0):

-

-        super().__init__()

-        self.dim = dim

-        self.input_resolution = input_resolution

-        self.depth = depth

-        self.use_checkpoint = use_checkpoint

-

-        # build blocks

-        self.blocks = nn.ModuleList([

-            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,

-                                 num_heads=num_heads, window_size=window_size,

-                                 shift_size=0 if (i % 2 == 0) else window_size // 2,

-                                 mlp_ratio=mlp_ratio,

-                                 qkv_bias=qkv_bias,

-                                 drop=drop, attn_drop=attn_drop,

-                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,

-                                 norm_layer=norm_layer,

-                                 pretrained_window_size=pretrained_window_size)

-            for i in range(depth)])

-

-        # patch merging layer

-        if downsample is not None:

-            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)

-        else:

-            self.downsample = None

-

-    def forward(self, x, x_size):

-        for blk in self.blocks:

-            if self.use_checkpoint:

-                x = checkpoint.checkpoint(blk, x, x_size)

-            else:

-                x = blk(x, x_size)

-        if self.downsample is not None:

-            x = self.downsample(x)

-        return x

-

-    def extra_repr(self) -> str:

-        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"

-

-    def flops(self):

-        flops = 0

-        for blk in self.blocks:

-            flops += blk.flops()

-        if self.downsample is not None:

-            flops += self.downsample.flops()

-        return flops

-

-    def _init_respostnorm(self):

-        for blk in self.blocks:

-            nn.init.constant_(blk.norm1.bias, 0)

-            nn.init.constant_(blk.norm1.weight, 0)

-            nn.init.constant_(blk.norm2.bias, 0)

-            nn.init.constant_(blk.norm2.weight, 0)

-            

-class PatchEmbed(nn.Module):

-    r""" Image to Patch Embedding

-    Args:

-        img_size (int): Image size.  Default: 224.

-        patch_size (int): Patch token size. Default: 4.

-        in_chans (int): Number of input image channels. Default: 3.

-        embed_dim (int): Number of linear projection output channels. Default: 96.

-        norm_layer (nn.Module, optional): Normalization layer. Default: None

-    """

-

-    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):

-        super().__init__()

-        img_size = to_2tuple(img_size)

-        patch_size = to_2tuple(patch_size)

-        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]

-        self.img_size = img_size

-        self.patch_size = patch_size

-        self.patches_resolution = patches_resolution

-        self.num_patches = patches_resolution[0] * patches_resolution[1]

-

-        self.in_chans = in_chans

-        self.embed_dim = embed_dim

-

-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

-        if norm_layer is not None:

-            self.norm = norm_layer(embed_dim)

-        else:

-            self.norm = None

-

-    def forward(self, x):

-        B, C, H, W = x.shape

-        # FIXME look at relaxing size constraints

-        # assert H == self.img_size[0] and W == self.img_size[1],

-        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."

-        x = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C

-        if self.norm is not None:

-            x = self.norm(x)

-        return x

-

-    def flops(self):

-        Ho, Wo = self.patches_resolution

-        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])

-        if self.norm is not None:

-            flops += Ho * Wo * self.embed_dim

-        return flops           

-

-class RSTB(nn.Module):

-    """Residual Swin Transformer Block (RSTB).

-

-    Args:

-        dim (int): Number of input channels.

-        input_resolution (tuple[int]): Input resolution.

-        depth (int): Number of blocks.

-        num_heads (int): Number of attention heads.

-        window_size (int): Local window size.

-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.

-        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True

-        drop (float, optional): Dropout rate. Default: 0.0

-        attn_drop (float, optional): Attention dropout rate. Default: 0.0

-        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0

-        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm

-        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None

-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.

-        img_size: Input image size.

-        patch_size: Patch size.

-        resi_connection: The convolutional block before residual connection.

-    """

-

-    def __init__(self, dim, input_resolution, depth, num_heads, window_size,

-                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0.,

-                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,

-                 img_size=224, patch_size=4, resi_connection='1conv'):

-        super(RSTB, self).__init__()

-

-        self.dim = dim

-        self.input_resolution = input_resolution

-

-        self.residual_group = BasicLayer(dim=dim,

-                                         input_resolution=input_resolution,

-                                         depth=depth,

-                                         num_heads=num_heads,

-                                         window_size=window_size,

-                                         mlp_ratio=mlp_ratio,

-                                         qkv_bias=qkv_bias, 

-                                         drop=drop, attn_drop=attn_drop,

-                                         drop_path=drop_path,

-                                         norm_layer=norm_layer,

-                                         downsample=downsample,

-                                         use_checkpoint=use_checkpoint)

-

-        if resi_connection == '1conv':

-            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)

-        elif resi_connection == '3conv':

-            # to save parameters and memory

-            self.conv = nn.Sequential(nn.Conv2d(dim, dim // 4, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2, inplace=True),

-                                      nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),

-                                      nn.LeakyReLU(negative_slope=0.2, inplace=True),

-                                      nn.Conv2d(dim // 4, dim, 3, 1, 1))

-

-        self.patch_embed = PatchEmbed(

-            img_size=img_size, patch_size=patch_size, in_chans=dim, embed_dim=dim,

-            norm_layer=None)

-

-        self.patch_unembed = PatchUnEmbed(

-            img_size=img_size, patch_size=patch_size, in_chans=dim, embed_dim=dim,

-            norm_layer=None)

-

-    def forward(self, x, x_size):

-        return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))) + x

-

-    def flops(self):

-        flops = 0

-        flops += self.residual_group.flops()

-        H, W = self.input_resolution

-        flops += H * W * self.dim * self.dim * 9

-        flops += self.patch_embed.flops()

-        flops += self.patch_unembed.flops()

-

-        return flops

-

-class PatchUnEmbed(nn.Module):

-    r""" Image to Patch Unembedding

-

-    Args:

-        img_size (int): Image size.  Default: 224.

-        patch_size (int): Patch token size. Default: 4.

-        in_chans (int): Number of input image channels. Default: 3.

-        embed_dim (int): Number of linear projection output channels. Default: 96.

-        norm_layer (nn.Module, optional): Normalization layer. Default: None

-    """

-

-    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):

-        super().__init__()

-        img_size = to_2tuple(img_size)

-        patch_size = to_2tuple(patch_size)

-        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]

-        self.img_size = img_size

-        self.patch_size = patch_size

-        self.patches_resolution = patches_resolution

-        self.num_patches = patches_resolution[0] * patches_resolution[1]

-

-        self.in_chans = in_chans

-        self.embed_dim = embed_dim

-

-    def forward(self, x, x_size):

-        B, HW, C = x.shape

-        x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C

-        return x

-

-    def flops(self):

-        flops = 0

-        return flops

-

-

-class Upsample(nn.Sequential):

-    """Upsample module.

-

-    Args:

-        scale (int): Scale factor. Supported scales: 2^n and 3.

-        num_feat (int): Channel number of intermediate features.

-    """

-

-    def __init__(self, scale, num_feat):

-        m = []

-        if (scale & (scale - 1)) == 0:  # scale = 2^n

-            for _ in range(int(math.log(scale, 2))):

-                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))

-                m.append(nn.PixelShuffle(2))

-        elif scale == 3:

-            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))

-            m.append(nn.PixelShuffle(3))

-        else:

-            raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')

-        super(Upsample, self).__init__(*m)

-        

-class Upsample_hf(nn.Sequential):

-    """Upsample module.

-

-    Args:

-        scale (int): Scale factor. Supported scales: 2^n and 3.

-        num_feat (int): Channel number of intermediate features.

-    """

-

-    def __init__(self, scale, num_feat):

-        m = []

-        if (scale & (scale - 1)) == 0:  # scale = 2^n

-            for _ in range(int(math.log(scale, 2))):

-                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))

-                m.append(nn.PixelShuffle(2))

-        elif scale == 3:

-            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))

-            m.append(nn.PixelShuffle(3))

-        else:

-            raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')

-        super(Upsample_hf, self).__init__(*m)        

-

-

-class UpsampleOneStep(nn.Sequential):

-    """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)

-       Used in lightweight SR to save parameters.

-

-    Args:

-        scale (int): Scale factor. Supported scales: 2^n and 3.

-        num_feat (int): Channel number of intermediate features.

-

-    """

-

-    def __init__(self, scale, num_feat, num_out_ch, input_resolution=None):

-        self.num_feat = num_feat

-        self.input_resolution = input_resolution

-        m = []

-        m.append(nn.Conv2d(num_feat, (scale ** 2) * num_out_ch, 3, 1, 1))

-        m.append(nn.PixelShuffle(scale))

-        super(UpsampleOneStep, self).__init__(*m)

-

-    def flops(self):

-        H, W = self.input_resolution

-        flops = H * W * self.num_feat * 3 * 9

-        return flops

-    

-    

-

-class Swin2SR(nn.Module):

-    r""" Swin2SR

-        A PyTorch impl of : `Swin2SR: SwinV2 Transformer for Compressed Image Super-Resolution and Restoration`.

-

-    Args:

-        img_size (int | tuple(int)): Input image size. Default 64

-        patch_size (int | tuple(int)): Patch size. Default: 1

-        in_chans (int): Number of input image channels. Default: 3

-        embed_dim (int): Patch embedding dimension. Default: 96

-        depths (tuple(int)): Depth of each Swin Transformer layer.

-        num_heads (tuple(int)): Number of attention heads in different layers.

-        window_size (int): Window size. Default: 7

-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4

-        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True

-        drop_rate (float): Dropout rate. Default: 0

-        attn_drop_rate (float): Attention dropout rate. Default: 0

-        drop_path_rate (float): Stochastic depth rate. Default: 0.1

-        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.

-        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False

-        patch_norm (bool): If True, add normalization after patch embedding. Default: True

-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False

-        upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction

-        img_range: Image range. 1. or 255.

-        upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None

-        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'

-    """

-

-    def __init__(self, img_size=64, patch_size=1, in_chans=3,

-                 embed_dim=96, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6],

-                 window_size=7, mlp_ratio=4., qkv_bias=True, 

-                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,

-                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,

-                 use_checkpoint=False, upscale=2, img_range=1., upsampler='', resi_connection='1conv',

-                 **kwargs):

-        super(Swin2SR, self).__init__()

-        num_in_ch = in_chans

-        num_out_ch = in_chans

-        num_feat = 64

-        self.img_range = img_range

-        if in_chans == 3:

-            rgb_mean = (0.4488, 0.4371, 0.4040)

-            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)

-        else:

-            self.mean = torch.zeros(1, 1, 1, 1)

-        self.upscale = upscale

-        self.upsampler = upsampler

-        self.window_size = window_size

-

-        #####################################################################################################

-        ################################### 1, shallow feature extraction ###################################

-        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)

-

-        #####################################################################################################

-        ################################### 2, deep feature extraction ######################################

-        self.num_layers = len(depths)

-        self.embed_dim = embed_dim

-        self.ape = ape

-        self.patch_norm = patch_norm

-        self.num_features = embed_dim

-        self.mlp_ratio = mlp_ratio

-

-        # split image into non-overlapping patches

-        self.patch_embed = PatchEmbed(

-            img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,

-            norm_layer=norm_layer if self.patch_norm else None)

-        num_patches = self.patch_embed.num_patches

-        patches_resolution = self.patch_embed.patches_resolution

-        self.patches_resolution = patches_resolution

-

-        # merge non-overlapping patches into image

-        self.patch_unembed = PatchUnEmbed(

-            img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,

-            norm_layer=norm_layer if self.patch_norm else None)

-

-        # absolute position embedding

-        if self.ape:

-            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))

-            trunc_normal_(self.absolute_pos_embed, std=.02)

-

-        self.pos_drop = nn.Dropout(p=drop_rate)

-

-        # stochastic depth

-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule

-

-        # build Residual Swin Transformer blocks (RSTB)

-        self.layers = nn.ModuleList()

-        for i_layer in range(self.num_layers):

-            layer = RSTB(dim=embed_dim,

-                         input_resolution=(patches_resolution[0],

-                                           patches_resolution[1]),

-                         depth=depths[i_layer],

-                         num_heads=num_heads[i_layer],

-                         window_size=window_size,

-                         mlp_ratio=self.mlp_ratio,

-                         qkv_bias=qkv_bias, 

-                         drop=drop_rate, attn_drop=attn_drop_rate,

-                         drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],  # no impact on SR results

-                         norm_layer=norm_layer,

-                         downsample=None,

-                         use_checkpoint=use_checkpoint,

-                         img_size=img_size,

-                         patch_size=patch_size,

-                         resi_connection=resi_connection

-

-                         )

-            self.layers.append(layer)

-            

-        if self.upsampler == 'pixelshuffle_hf':

-            self.layers_hf = nn.ModuleList()

-            for i_layer in range(self.num_layers):

-                layer = RSTB(dim=embed_dim,

-                             input_resolution=(patches_resolution[0],

-                                               patches_resolution[1]),

-                             depth=depths[i_layer],

-                             num_heads=num_heads[i_layer],

-                             window_size=window_size,

-                             mlp_ratio=self.mlp_ratio,

-                             qkv_bias=qkv_bias, 

-                             drop=drop_rate, attn_drop=attn_drop_rate,

-                             drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],  # no impact on SR results

-                             norm_layer=norm_layer,

-                             downsample=None,

-                             use_checkpoint=use_checkpoint,

-                             img_size=img_size,

-                             patch_size=patch_size,

-                             resi_connection=resi_connection

-

-                             )

-                self.layers_hf.append(layer)

-                        

-        self.norm = norm_layer(self.num_features)

-

-        # build the last conv layer in deep feature extraction

-        if resi_connection == '1conv':

-            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)

-        elif resi_connection == '3conv':

-            # to save parameters and memory

-            self.conv_after_body = nn.Sequential(nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),

-                                                 nn.LeakyReLU(negative_slope=0.2, inplace=True),

-                                                 nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),

-                                                 nn.LeakyReLU(negative_slope=0.2, inplace=True),

-                                                 nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1))

-

-        #####################################################################################################

-        ################################ 3, high quality image reconstruction ################################

-        if self.upsampler == 'pixelshuffle':

-            # for classical SR

-            self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),

-                                                      nn.LeakyReLU(inplace=True))

-            self.upsample = Upsample(upscale, num_feat)

-            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

-        elif self.upsampler == 'pixelshuffle_aux':

-            self.conv_bicubic = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)

-            self.conv_before_upsample = nn.Sequential(

-                nn.Conv2d(embed_dim, num_feat, 3, 1, 1),

-                nn.LeakyReLU(inplace=True))

-            self.conv_aux = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

-            self.conv_after_aux = nn.Sequential(

-                nn.Conv2d(3, num_feat, 3, 1, 1),

-                nn.LeakyReLU(inplace=True))            

-            self.upsample = Upsample(upscale, num_feat)

-            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

-            

-        elif self.upsampler == 'pixelshuffle_hf':

-            self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),

-                                                      nn.LeakyReLU(inplace=True))

-            self.upsample = Upsample(upscale, num_feat)

-            self.upsample_hf = Upsample_hf(upscale, num_feat)

-            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

-            self.conv_first_hf = nn.Sequential(nn.Conv2d(num_feat, embed_dim, 3, 1, 1),

-                                                      nn.LeakyReLU(inplace=True))

-            self.conv_after_body_hf = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)

-            self.conv_before_upsample_hf = nn.Sequential(

-                nn.Conv2d(embed_dim, num_feat, 3, 1, 1),

-                nn.LeakyReLU(inplace=True))

-            self.conv_last_hf = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

-            

-        elif self.upsampler == 'pixelshuffledirect':

-            # for lightweight SR (to save parameters)

-            self.upsample = UpsampleOneStep(upscale, embed_dim, num_out_ch,

-                                            (patches_resolution[0], patches_resolution[1]))

-        elif self.upsampler == 'nearest+conv':

-            # for real-world SR (less artifacts)

-            assert self.upscale == 4, 'only support x4 now.'

-            self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),

-                                                      nn.LeakyReLU(inplace=True))

-            self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)

-            self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)

-            self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)

-            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

-            self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

-        else:

-            # for image denoising and JPEG compression artifact reduction

-            self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1)

-

-        self.apply(self._init_weights)

-

-    def _init_weights(self, m):

-        if isinstance(m, nn.Linear):

-            trunc_normal_(m.weight, std=.02)

-            if isinstance(m, nn.Linear) and m.bias is not None:

-                nn.init.constant_(m.bias, 0)

-        elif isinstance(m, nn.LayerNorm):

-            nn.init.constant_(m.bias, 0)

-            nn.init.constant_(m.weight, 1.0)

-

-    @torch.jit.ignore

-    def no_weight_decay(self):

-        return {'absolute_pos_embed'}

-

-    @torch.jit.ignore

-    def no_weight_decay_keywords(self):

-        return {'relative_position_bias_table'}

-

-    def check_image_size(self, x):

-        _, _, h, w = x.size()

-        mod_pad_h = (self.window_size - h % self.window_size) % self.window_size

-        mod_pad_w = (self.window_size - w % self.window_size) % self.window_size

-        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect')

-        return x

-

-    def forward_features(self, x):

-        x_size = (x.shape[2], x.shape[3])

-        x = self.patch_embed(x)

-        if self.ape:

-            x = x + self.absolute_pos_embed

-        x = self.pos_drop(x)

-

-        for layer in self.layers:

-            x = layer(x, x_size)

-

-        x = self.norm(x)  # B L C

-        x = self.patch_unembed(x, x_size)

-

-        return x

-    

-    def forward_features_hf(self, x):

-        x_size = (x.shape[2], x.shape[3])

-        x = self.patch_embed(x)

-        if self.ape:

-            x = x + self.absolute_pos_embed

-        x = self.pos_drop(x)

-

-        for layer in self.layers_hf:

-            x = layer(x, x_size)

-

-        x = self.norm(x)  # B L C

-        x = self.patch_unembed(x, x_size)

-

-        return x    

-

-    def forward(self, x):

-        H, W = x.shape[2:]

-        x = self.check_image_size(x)

-

-        self.mean = self.mean.type_as(x)

-        x = (x - self.mean) * self.img_range

-

-        if self.upsampler == 'pixelshuffle':

-            # for classical SR

-            x = self.conv_first(x)

-            x = self.conv_after_body(self.forward_features(x)) + x

-            x = self.conv_before_upsample(x)

-            x = self.conv_last(self.upsample(x))

-        elif self.upsampler == 'pixelshuffle_aux':

-            bicubic = F.interpolate(x, size=(H * self.upscale, W * self.upscale), mode='bicubic', align_corners=False)

-            bicubic = self.conv_bicubic(bicubic)

-            x = self.conv_first(x)

-            x = self.conv_after_body(self.forward_features(x)) + x

-            x = self.conv_before_upsample(x)

-            aux = self.conv_aux(x) # b, 3, LR_H, LR_W

-            x = self.conv_after_aux(aux)

-            x = self.upsample(x)[:, :, :H * self.upscale, :W * self.upscale] + bicubic[:, :, :H * self.upscale, :W * self.upscale]

-            x = self.conv_last(x)

-            aux = aux / self.img_range + self.mean

-        elif self.upsampler == 'pixelshuffle_hf':

-            # for classical SR with HF

-            x = self.conv_first(x)

-            x = self.conv_after_body(self.forward_features(x)) + x

-            x_before = self.conv_before_upsample(x)

-            x_out = self.conv_last(self.upsample(x_before))

-            

-            x_hf = self.conv_first_hf(x_before)

-            x_hf = self.conv_after_body_hf(self.forward_features_hf(x_hf)) + x_hf

-            x_hf = self.conv_before_upsample_hf(x_hf)

-            x_hf = self.conv_last_hf(self.upsample_hf(x_hf))

-            x = x_out + x_hf

-            x_hf = x_hf / self.img_range + self.mean

-

-        elif self.upsampler == 'pixelshuffledirect':

-            # for lightweight SR

-            x = self.conv_first(x)

-            x = self.conv_after_body(self.forward_features(x)) + x

-            x = self.upsample(x)

-        elif self.upsampler == 'nearest+conv':

-            # for real-world SR

-            x = self.conv_first(x)

-            x = self.conv_after_body(self.forward_features(x)) + x

-            x = self.conv_before_upsample(x)

-            x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))

-            x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))

-            x = self.conv_last(self.lrelu(self.conv_hr(x)))

-        else:

-            # for image denoising and JPEG compression artifact reduction

-            x_first = self.conv_first(x)

-            res = self.conv_after_body(self.forward_features(x_first)) + x_first

-            x = x + self.conv_last(res)

-        

-        x = x / self.img_range + self.mean

-        if self.upsampler == "pixelshuffle_aux":

-            return x[:, :, :H*self.upscale, :W*self.upscale], aux

-        

-        elif self.upsampler == "pixelshuffle_hf":

-            x_out = x_out / self.img_range + self.mean

-            return x_out[:, :, :H*self.upscale, :W*self.upscale], x[:, :, :H*self.upscale, :W*self.upscale], x_hf[:, :, :H*self.upscale, :W*self.upscale]

-        

-        else:

-            return x[:, :, :H*self.upscale, :W*self.upscale]

-

-    def flops(self):

-        flops = 0

-        H, W = self.patches_resolution

-        flops += H * W * 3 * self.embed_dim * 9

-        flops += self.patch_embed.flops()

-        for i, layer in enumerate(self.layers):

-            flops += layer.flops()

-        flops += H * W * 3 * self.embed_dim * self.embed_dim

-        flops += self.upsample.flops()

-        return flops

-

-

-if __name__ == '__main__':

-    upscale = 4

-    window_size = 8

-    height = (1024 // upscale // window_size + 1) * window_size

-    width = (720 // upscale // window_size + 1) * window_size

-    model = Swin2SR(upscale=2, img_size=(height, width),

-                   window_size=window_size, img_range=1., depths=[6, 6, 6, 6],

-                   embed_dim=60, num_heads=[6, 6, 6, 6], mlp_ratio=2, upsampler='pixelshuffledirect')

-    print(model)

-    print(height, width, model.flops() / 1e9)

-

-    x = torch.randn((1, 3, height, width))

-    x = model(x)

-    print(x.shape)
\ No newline at end of file
diff --git a/modules/ui.py b/modules/ui.py
index 2eb0b684..3acb9b48 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -28,7 +28,6 @@ import modules.codeformer_model
 import modules.generation_parameters_copypaste as parameters_copypaste

 import modules.gfpgan_model

 import modules.hypernetworks.ui

-import modules.ldsr_model

 import modules.scripts

 import modules.shared as shared

 import modules.styles

diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
index 42667941..b487ac25 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -78,6 +78,12 @@ def extension_table():
     """

 

     for ext in extensions.extensions:

+        remote = ""

+        if ext.is_builtin:

+            remote = "built-in"

+        elif ext.remote:

+            remote = f"""<a href="{html.escape(ext.remote or '')}" target="_blank">{html.escape("built-in" if ext.is_builtin else ext.remote or '')}</a>"""

+

         if ext.can_update:

             ext_status = f"""<label><input class="gr-check-radio gr-checkbox" name="update_{html.escape(ext.name)}" checked="checked" type="checkbox">{html.escape(ext.status)}</label>"""

         else:

@@ -86,7 +92,7 @@ def extension_table():
         code += f"""

             <tr>

                 <td><label><input class="gr-check-radio gr-checkbox" name="enable_{html.escape(ext.name)}" type="checkbox" {'checked="checked"' if ext.enabled else ''}>{html.escape(ext.name)}</label></td>

-                <td><a href="{html.escape(ext.remote or '')}" target="_blank">{html.escape(ext.remote or '')}</a></td>

+                <td>{remote}</td>

                 <td{' class="extension_status"' if ext.remote is not None else ''}>{ext_status}</td>

             </tr>

     """