update swinir

extensions-builtin/SwinIR/scripts/swinir_model.py

@@ -1,11 +1,10 @@
-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 rich import print, progress # pylint: disable=redefined-builtin
 
 from modules import modelloader, devices, script_callbacks, shared
 from modules.shared import cmd_opts, opts, state
@@ -59,17 +58,17 @@ class UpscalerSwinIR(Upscaler):
             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",
+                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:
@@ -88,7 +87,8 @@ class UpscalerSwinIR(Upscaler):
             )
             params = "params_ema"
 
-        pretrained_model = torch.load(filename)
+        with progress.open(filename, 'rb', description=f'Loading weights: [cyan]{filename}', auto_refresh=True) as f:
+            pretrained_model = torch.load(filename)
         if params is not None:
             model.load_state_dict(pretrained_model[params], strict=True)
         else: