Merge branch 'dev' into text-drag-fix

3 files changed, 36 insertions, 37 deletions

diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/extensions-builtin/SwinIR/scripts/swinir_model.py
index e8783bca..1c7bf325 100644
--- a/extensions-builtin/SwinIR/scripts/swinir_model.py
+++ b/extensions-builtin/SwinIR/scripts/swinir_model.py
@@ -1,4 +1,3 @@
-import contextlib
 import os
 
 import numpy as np
@@ -8,7 +7,7 @@ from basicsr.utils.download_util import load_file_from_url
 from tqdm import tqdm
 
 from modules import modelloader, devices, script_callbacks, shared
-from modules.shared import cmd_opts, opts, state
+from modules.shared import opts, state
 from swinir_model_arch import SwinIR as net
 from swinir_model_arch_v2 import Swin2SR as net2
 from modules.upscaler import Upscaler, UpscalerData
@@ -45,14 +44,14 @@ class UpscalerSwinIR(Upscaler):
         img = upscale(img, model)
         try:
             torch.cuda.empty_cache()
-        except:
+        except Exception:
             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)
+            filename = load_file_from_url(url=path, model_dir=self.model_download_path, file_name=dl_name, progress=True)
         else:
             filename = path
         if filename is None or not os.path.exists(filename):
@@ -151,7 +150,7 @@ def inference(img, model, tile, tile_overlap, window_size, scale):
             for w_idx in w_idx_list:
                 if state.interrupted or state.skipped:
                     break
-                
+
                 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)
diff --git a/extensions-builtin/SwinIR/swinir_model_arch.py b/extensions-builtin/SwinIR/swinir_model_arch.py
index 863f42db..93b93274 100644
--- a/extensions-builtin/SwinIR/swinir_model_arch.py
+++ b/extensions-builtin/SwinIR/swinir_model_arch.py
@@ -644,7 +644,7 @@ class SwinIR(nn.Module):
     """
 
     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],
+                 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,
@@ -805,7 +805,7 @@ class SwinIR(nn.Module):
     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
 
@@ -844,7 +844,7 @@ class SwinIR(nn.Module):
         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):
+        for layer in self.layers:
             flops += layer.flops()
         flops += H * W * 3 * self.embed_dim * self.embed_dim
         flops += self.upsample.flops()
diff --git a/extensions-builtin/SwinIR/swinir_model_arch_v2.py b/extensions-builtin/SwinIR/swinir_model_arch_v2.py
index 0e28ae6e..dad22cca 100644
--- a/extensions-builtin/SwinIR/swinir_model_arch_v2.py
+++ b/extensions-builtin/SwinIR/swinir_model_arch_v2.py
@@ -74,7 +74,7 @@ class WindowAttention(nn.Module):
     """

 

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

-                 pretrained_window_size=[0, 0]):

+                 pretrained_window_size=(0, 0)):

 

         super().__init__()

         self.dim = dim

@@ -241,7 +241,7 @@ class SwinTransformerBlock(nn.Module):
             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

@@ -263,7 +263,7 @@ class SwinTransformerBlock(nn.Module):
         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        

+        return attn_mask

 

     def forward(self, x, x_size):

         H, W = x_size

@@ -288,7 +288,7 @@ class SwinTransformerBlock(nn.Module):
             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

@@ -369,7 +369,7 @@ class PatchMerging(nn.Module):
         H, W = self.input_resolution

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

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

-        return flops    

+        return flops

 

 class BasicLayer(nn.Module):

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

@@ -447,7 +447,7 @@ class BasicLayer(nn.Module):
             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:

@@ -492,7 +492,7 @@ class PatchEmbed(nn.Module):
         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           

+        return flops

 

 class RSTB(nn.Module):

     """Residual Swin Transformer Block (RSTB).

@@ -531,7 +531,7 @@ class RSTB(nn.Module):
                                          num_heads=num_heads,

                                          window_size=window_size,

                                          mlp_ratio=mlp_ratio,

-                                         qkv_bias=qkv_bias, 

+                                         qkv_bias=qkv_bias,

                                          drop=drop, attn_drop=attn_drop,

                                          drop_path=drop_path,

                                          norm_layer=norm_layer,

@@ -622,7 +622,7 @@ class Upsample(nn.Sequential):
         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.

 

@@ -642,7 +642,7 @@ class Upsample_hf(nn.Sequential):
             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)        

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

 

 

 class UpsampleOneStep(nn.Sequential):

@@ -667,8 +667,8 @@ class UpsampleOneStep(nn.Sequential):
         H, W = self.input_resolution

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

         return flops

-    

-    

+

+

 

 class Swin2SR(nn.Module):

     r""" Swin2SR

@@ -698,8 +698,8 @@ class Swin2SR(nn.Module):
     """

 

     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, 

+                 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',

@@ -764,7 +764,7 @@ class Swin2SR(nn.Module):
                          num_heads=num_heads[i_layer],

                          window_size=window_size,

                          mlp_ratio=self.mlp_ratio,

-                         qkv_bias=qkv_bias, 

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

@@ -776,7 +776,7 @@ class Swin2SR(nn.Module):
 

                          )

             self.layers.append(layer)

-            

+

         if self.upsampler == 'pixelshuffle_hf':

             self.layers_hf = nn.ModuleList()

             for i_layer in range(self.num_layers):

@@ -787,7 +787,7 @@ class Swin2SR(nn.Module):
                              num_heads=num_heads[i_layer],

                              window_size=window_size,

                              mlp_ratio=self.mlp_ratio,

-                             qkv_bias=qkv_bias, 

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

@@ -799,7 +799,7 @@ class Swin2SR(nn.Module):
 

                              )

                 self.layers_hf.append(layer)

-                        

+

         self.norm = norm_layer(self.num_features)

 

         # build the last conv layer in deep feature extraction

@@ -829,10 +829,10 @@ class Swin2SR(nn.Module):
             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))            

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

@@ -846,7 +846,7 @@ class Swin2SR(nn.Module):
                 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,

@@ -905,7 +905,7 @@ class Swin2SR(nn.Module):
         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)

@@ -919,7 +919,7 @@ class Swin2SR(nn.Module):
         x = self.norm(x)  # B L C

         x = self.patch_unembed(x, x_size)

 

-        return x    

+        return x

 

     def forward(self, x):

         H, W = x.shape[2:]

@@ -951,7 +951,7 @@ class Swin2SR(nn.Module):
             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)

@@ -977,15 +977,15 @@ class Swin2SR(nn.Module):
             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]

 

@@ -994,7 +994,7 @@ class Swin2SR(nn.Module):
         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):

+        for layer in self.layers:

             flops += layer.flops()

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

         flops += self.upsample.flops()

@@ -1014,4 +1014,4 @@ if __name__ == '__main__':
 

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

     x = model(x)

-    print(x.shape)
\ No newline at end of file
+    print(x.shape)