Merge pull request #2 from JC-Array/master

resolve merge conflicts

14 files changed, 1171 insertions, 60 deletions

diff --git a/modules/devices.py b/modules/devices.py
index 0158b11f..03ef58f1 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -36,6 +36,7 @@ errors.run(enable_tf32, "Enabling TF32")
 
 device = device_gfpgan = device_bsrgan = device_esrgan = device_scunet = device_codeformer = get_optimal_device()
 dtype = torch.float16
+dtype_vae = torch.float16
 
 def randn(seed, shape):
     # Pytorch currently doesn't handle setting randomness correctly when the metal backend is used.
@@ -59,9 +60,12 @@ def randn_without_seed(shape):
     return torch.randn(shape, device=device)
 
 
-def autocast():
+def autocast(disable=False):
     from modules import shared
 
+    if disable:
+        return contextlib.nullcontext()
+
     if dtype == torch.float32 or shared.cmd_opts.precision == "full":
         return contextlib.nullcontext()
 
diff --git a/modules/processing.py b/modules/processing.py
index 94d2dd62..698b3069 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -207,7 +207,7 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
     # enables the generation of additional tensors with noise that the sampler will use during its processing.

     # Using those pre-generated tensors instead of simple torch.randn allows a batch with seeds [100, 101] to

     # produce the same images as with two batches [100], [101].

-    if p is not None and p.sampler is not None and len(seeds) > 1 and opts.enable_batch_seeds:

+    if p is not None and p.sampler is not None and (len(seeds) > 1 and opts.enable_batch_seeds or opts.eta_noise_seed_delta > 0):

         sampler_noises = [[] for _ in range(p.sampler.number_of_needed_noises(p))]

     else:

         sampler_noises = None

@@ -247,6 +247,9 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
         if sampler_noises is not None:

             cnt = p.sampler.number_of_needed_noises(p)

 

+            if opts.eta_noise_seed_delta > 0:

+                torch.manual_seed(seed + opts.eta_noise_seed_delta)

+

             for j in range(cnt):

                 sampler_noises[j].append(devices.randn_without_seed(tuple(noise_shape)))

 

@@ -259,6 +262,13 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
     return x

 

 

+def decode_first_stage(model, x):

+    with devices.autocast(disable=x.dtype == devices.dtype_vae):

+        x = model.decode_first_stage(x)

+

+    return x

+

+

 def get_fixed_seed(seed):

     if seed is None or seed == '' or seed == -1:

         return int(random.randrange(4294967294))

@@ -294,6 +304,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration
         "Denoising strength": getattr(p, 'denoising_strength', None),

         "Eta": (None if p.sampler is None or p.sampler.eta == p.sampler.default_eta else p.sampler.eta),

         "Clip skip": None if clip_skip <= 1 else clip_skip,

+        "ENSD": None if opts.eta_noise_seed_delta == 0 else opts.eta_noise_seed_delta,

     }

 

     generation_params.update(p.extra_generation_params)

@@ -398,9 +409,8 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
                 # use the image collected previously in sampler loop

                 samples_ddim = shared.state.current_latent

 

-            samples_ddim = samples_ddim.to(devices.dtype)

-

-            x_samples_ddim = p.sd_model.decode_first_stage(samples_ddim)

+            samples_ddim = samples_ddim.to(devices.dtype_vae)

+            x_samples_ddim = decode_first_stage(p.sd_model, samples_ddim)

             x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)

 

             del samples_ddim

@@ -533,7 +543,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         if self.scale_latent:

             samples = torch.nn.functional.interpolate(samples, size=(self.height // opt_f, self.width // opt_f), mode="bilinear")

         else:

-            decoded_samples = self.sd_model.decode_first_stage(samples)

+            decoded_samples = decode_first_stage(self.sd_model, samples)

 

             if opts.upscaler_for_img2img is None or opts.upscaler_for_img2img == "None":

                 decoded_samples = torch.nn.functional.interpolate(decoded_samples, size=(self.height, self.width), mode="bilinear")

diff --git a/modules/safe.py b/modules/safe.py
index 4d06f2a5..05917463 100644
--- a/modules/safe.py
+++ b/modules/safe.py
@@ -12,6 +12,10 @@ import _codecs
 import zipfile

 

 

+# PyTorch 1.13 and later have _TypedStorage renamed to TypedStorage

+TypedStorage = torch.storage.TypedStorage if hasattr(torch.storage, 'TypedStorage') else torch.storage._TypedStorage

+

+

 def encode(*args):

     out = _codecs.encode(*args)

     return out

@@ -20,7 +24,7 @@ def encode(*args):
 class RestrictedUnpickler(pickle.Unpickler):

     def persistent_load(self, saved_id):

         assert saved_id[0] == 'storage'

-        return torch.storage._TypedStorage()

+        return TypedStorage()

 

     def find_class(self, module, name):

         if module == 'collections' and name == 'OrderedDict':

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 437acce4..827bf304 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -23,7 +23,7 @@ def apply_optimizations():
 

     ldm.modules.diffusionmodules.model.nonlinearity = silu

 

-    if cmd_opts.force_enable_xformers or (cmd_opts.xformers and shared.xformers_available and torch.version.cuda and torch.cuda.get_device_capability(shared.device) == (8, 6)):

+    if cmd_opts.force_enable_xformers or (cmd_opts.xformers and shared.xformers_available and torch.version.cuda and (6, 0) <= torch.cuda.get_device_capability(shared.device) <= (8, 6)):

         print("Applying xformers cross attention optimization.")

         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.xformers_attention_forward

         ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.xformers_attnblock_forward

@@ -43,10 +43,7 @@ def undo_optimizations():
 

 

 def get_target_prompt_token_count(token_count):

-    if token_count < 75:

-        return 75

-

-    return math.ceil(token_count / 10) * 10

+    return math.ceil(max(token_count, 1) / 75) * 75

 

 

 class StableDiffusionModelHijack:

@@ -127,7 +124,6 @@ class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
                 self.token_mults[ident] = mult

 

     def tokenize_line(self, line, used_custom_terms, hijack_comments):

-        id_start = self.wrapped.tokenizer.bos_token_id

         id_end = self.wrapped.tokenizer.eos_token_id

 

         if opts.enable_emphasis:

@@ -154,7 +150,13 @@ class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
                     i += 1

                 else:

                     emb_len = int(embedding.vec.shape[0])

-                    fixes.append((len(remade_tokens), embedding))

+                    iteration = len(remade_tokens) // 75

+                    if (len(remade_tokens) + emb_len) // 75 != iteration:

+                        rem = (75 * (iteration + 1) - len(remade_tokens))

+                        remade_tokens += [id_end] * rem

+                        multipliers += [1.0] * rem

+                        iteration += 1

+                    fixes.append((iteration, (len(remade_tokens) % 75, embedding)))

                     remade_tokens += [0] * emb_len

                     multipliers += [weight] * emb_len

                     used_custom_terms.append((embedding.name, embedding.checksum()))

@@ -162,10 +164,10 @@ class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
 

         token_count = len(remade_tokens)

         prompt_target_length = get_target_prompt_token_count(token_count)

-        tokens_to_add = prompt_target_length - len(remade_tokens) + 1

+        tokens_to_add = prompt_target_length - len(remade_tokens)

 

-        remade_tokens = [id_start] + remade_tokens + [id_end] * tokens_to_add

-        multipliers = [1.0] + multipliers + [1.0] * tokens_to_add

+        remade_tokens = remade_tokens + [id_end] * tokens_to_add

+        multipliers = multipliers + [1.0] * tokens_to_add

 

         return remade_tokens, fixes, multipliers, token_count

 

@@ -260,29 +262,55 @@ class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
             hijack_fixes.append(fixes)

             batch_multipliers.append(multipliers)

         return batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count

-

+    

     def forward(self, text):

-

-        if opts.use_old_emphasis_implementation:

+        use_old = opts.use_old_emphasis_implementation

+        if use_old:

             batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count = self.process_text_old(text)

         else:

             batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count = self.process_text(text)

 

-        self.hijack.fixes = hijack_fixes

         self.hijack.comments += hijack_comments

 

         if len(used_custom_terms) > 0:

             self.hijack.comments.append("Used embeddings: " + ", ".join([f'{word} [{checksum}]' for word, checksum in used_custom_terms]))

+        

+        if use_old:

+            self.hijack.fixes = hijack_fixes

+            return self.process_tokens(remade_batch_tokens, batch_multipliers)

+        

+        z = None

+        i = 0

+        while max(map(len, remade_batch_tokens)) != 0:

+            rem_tokens = [x[75:] for x in remade_batch_tokens]

+            rem_multipliers = [x[75:] for x in batch_multipliers]

+            

+            self.hijack.fixes = []

+            for unfiltered in hijack_fixes:

+                fixes = []

+                for fix in unfiltered:

+                    if fix[0] == i:

+                        fixes.append(fix[1])

+                self.hijack.fixes.append(fixes)

+            

+            z1 = self.process_tokens([x[:75] for x in remade_batch_tokens], [x[:75] for x in batch_multipliers])

+            z = z1 if z is None else torch.cat((z, z1), axis=-2)

+            

+            remade_batch_tokens = rem_tokens

+            batch_multipliers = rem_multipliers

+            i += 1

+        

+        return z

+        

+    

+    def process_tokens(self, remade_batch_tokens, batch_multipliers):

+        if not opts.use_old_emphasis_implementation:

+            remade_batch_tokens = [[self.wrapped.tokenizer.bos_token_id] + x[:75] + [self.wrapped.tokenizer.eos_token_id] for x in remade_batch_tokens]

+            batch_multipliers = [[1.0] + x[:75] + [1.0] for x in batch_multipliers]

+        

+        tokens = torch.asarray(remade_batch_tokens).to(device)

+        outputs = self.wrapped.transformer(input_ids=tokens, output_hidden_states=-opts.CLIP_stop_at_last_layers)

 

-        target_token_count = get_target_prompt_token_count(token_count) + 2

-

-        position_ids_array = [min(x, 75) for x in range(target_token_count-1)] + [76]

-        position_ids = torch.asarray(position_ids_array, device=devices.device).expand((1, -1))

-

-        remade_batch_tokens_of_same_length = [x + [self.wrapped.tokenizer.eos_token_id] * (target_token_count - len(x)) for x in remade_batch_tokens]

-        tokens = torch.asarray(remade_batch_tokens_of_same_length).to(device)

-

-        outputs = self.wrapped.transformer(input_ids=tokens, position_ids=position_ids, output_hidden_states=-opts.CLIP_stop_at_last_layers)

         if opts.CLIP_stop_at_last_layers > 1:

             z = outputs.hidden_states[-opts.CLIP_stop_at_last_layers]

             z = self.wrapped.transformer.text_model.final_layer_norm(z)

@@ -290,7 +318,7 @@ class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
             z = outputs.last_hidden_state

 

         # restoring original mean is likely not correct, but it seems to work well to prevent artifacts that happen otherwise

-        batch_multipliers_of_same_length = [x + [1.0] * (target_token_count - len(x)) for x in batch_multipliers]

+        batch_multipliers_of_same_length = [x + [1.0] * (75 - len(x)) for x in batch_multipliers]

         batch_multipliers = torch.asarray(batch_multipliers_of_same_length).to(device)

         original_mean = z.mean()

         z *= batch_multipliers.reshape(batch_multipliers.shape + (1,)).expand(z.shape)

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 634fb4b2..18408e62 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -13,8 +13,6 @@ from modules import shared
 if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:

     try:

         import xformers.ops

-        import functorch

-        xformers._is_functorch_available = True

         shared.xformers_available = True

     except Exception:

         print("Cannot import xformers", file=sys.stderr)

diff --git a/modules/sd_models.py b/modules/sd_models.py
index e63d3c29..0a55b4c3 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -149,8 +149,13 @@ def load_model_weights(model, checkpoint_info):
         model.half()

 

     devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16

+    devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16

 

     vae_file = os.path.splitext(checkpoint_file)[0] + ".vae.pt"

+

+    if not os.path.exists(vae_file) and shared.cmd_opts.vae_path is not None:

+        vae_file = shared.cmd_opts.vae_path

+

     if os.path.exists(vae_file):

         print(f"Loading VAE weights from: {vae_file}")

         vae_ckpt = torch.load(vae_file, map_location="cpu")

@@ -158,6 +163,8 @@ def load_model_weights(model, checkpoint_info):
 

         model.first_stage_model.load_state_dict(vae_dict)

 

+    model.first_stage_model.to(devices.dtype_vae)

+

     model.sd_model_hash = sd_model_hash

     model.sd_model_checkpoint = checkpoint_file

     model.sd_checkpoint_info = checkpoint_info

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 6e743f7e..d168b938 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -7,7 +7,7 @@ import inspect
 import k_diffusion.sampling

 import ldm.models.diffusion.ddim

 import ldm.models.diffusion.plms

-from modules import prompt_parser

+from modules import prompt_parser, devices, processing

 

 from modules.shared import opts, cmd_opts, state

 import modules.shared as shared

@@ -83,7 +83,7 @@ def setup_img2img_steps(p, steps=None):
 

 

 def sample_to_image(samples):

-    x_sample = shared.sd_model.decode_first_stage(samples[0:1].type(shared.sd_model.dtype))[0]

+    x_sample = processing.decode_first_stage(shared.sd_model, samples[0:1])[0]

     x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)

     x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)

     x_sample = x_sample.astype(np.uint8)

diff --git a/modules/shared.py b/modules/shared.py
index 2e307809..99a0264c 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -25,6 +25,7 @@ parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to director
 parser.add_argument("--gfpgan-dir", type=str, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN'))

 parser.add_argument("--gfpgan-model", type=str, help="GFPGAN model file name", default=None)

 parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats")

+parser.add_argument("--no-half-vae", action='store_true', help="do not switch the VAE model to 16-bit floats")

 parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")

 parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")

 parser.add_argument("--embeddings-dir", type=str, default=os.path.join(script_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")

@@ -65,6 +66,7 @@ parser.add_argument("--autolaunch", action='store_true', help="open the webui UR
 parser.add_argument("--use-textbox-seed", action='store_true', help="use textbox for seeds in UI (no up/down, but possible to input long seeds)", default=False)

 parser.add_argument("--disable-console-progressbars", action='store_true', help="do not output progressbars to console", default=False)

 parser.add_argument("--enable-console-prompts", action='store_true', help="print prompts to console when generating with txt2img and img2img", default=False)

+parser.add_argument('--vae-path', type=str, help='Path to Variational Autoencoders model', default=None)

 parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=False)

 

 

@@ -171,6 +173,7 @@ options_templates.update(options_section(('saving-images', "Saving images/grids"
 

     "use_original_name_batch": OptionInfo(False, "Use original name for output filename during batch process in extras tab"),

     "save_selected_only": OptionInfo(True, "When using 'Save' button, only save a single selected image"),

+    "do_not_add_watermark": OptionInfo(False, "Do not add watermark to images"),

 }))

 

 options_templates.update(options_section(('saving-paths', "Paths for saving"), {

@@ -259,6 +262,7 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters"
     's_churn': OptionInfo(0.0, "sigma churn", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

     's_tmin':  OptionInfo(0.0, "sigma tmin",  gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

     's_noise': OptionInfo(1.0, "sigma noise", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

+    'eta_noise_seed_delta': OptionInfo(0, "Eta noise seed delta", gr.Number, {"precision": 0}),

 }))

 

 if cmd_opts.deepdanbooru:

diff --git a/modules/swinir_model.py b/modules/swinir_model.py
index fbd11f84..baa02e3d 100644
--- a/modules/swinir_model.py
+++ b/modules/swinir_model.py
@@ -10,6 +10,7 @@ from tqdm import tqdm
 from modules import modelloader
 from modules.shared import cmd_opts, opts, device
 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
 
 precision_scope = (
@@ -57,22 +58,42 @@ class UpscalerSwinIR(Upscaler):
             filename = path
         if filename is None or not os.path.exists(filename):
             return None
-        model = net(
+        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, 6, 6, 6],
-            embed_dim=240,
-            num_heads=[8, 8, 8, 8, 8, 8, 8, 8, 8],
+            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="3conv",
-        )
+            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)
-        model.load_state_dict(pretrained_model["params_ema"], strict=True)
+        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
diff --git a/modules/swinir_model_arch_v2.py b/modules/swinir_model_arch_v2.py
new file mode 100644
index 00000000..0e28ae6e
--- /dev/null
+++ b/modules/swinir_model_arch_v2.py
@@ -0,0 +1,1017 @@
+# -----------------------------------------------------------------------------------

+# 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

+