feat(api): add params for SwinIR large models

ssube · Apr 11, 2023 · 7f6da8e · 7f6da8e
1 parent 6570d19
commit 7f6da8e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/api/onnx_web/convert/__main__.py b/api/onnx_web/convert/__main__.py
@@ -118,7 +118,7 @@
         },
         {
             "model": "swinir",
-            "name": "upscaling-swinir-real-x4",
+            "name": "upscaling-swinir-real-large-x4",
             "source": "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth",
             "scale": 4,
         },

diff --git a/api/onnx_web/convert/upscaling/swinir.py b/api/onnx_web/convert/upscaling/swinir.py
@@ -48,9 +48,16 @@ def convert_upscaling_swinir(
         params["num_heads"] = [6, 6, 6, 6]
         params["upsampler"] = "pixelshuffledirect"
     elif "real" in name:
-        # TODO: add params for large model
-        logger.debug("using SwinIR real params")
-        params["upsampler"] = "nearest+conv"
+        if "large" in name:
+            logger.debug("using SwinIR real large params")
+            params["depths"] = [6] * 9
+            params["embed_dim"] = 240
+            params["num_heads"] = [8] * 9
+            params["upsampler"] = "nearest+conv"
+            params["resi_connection"] = "3conv"
+        else:
+            logger.debug("using SwinIR real params")
+            params["upsampler"] = "nearest+conv"
     elif "gray_dn" in name:
         params["img_size"] = (128, 128)
         params["in_chans"] = 1
@@ -59,11 +66,13 @@ def convert_upscaling_swinir(
         params["img_size"] = (128, 128)
         params["upsampler"] = ""
     elif "gray_jpeg" in name:
+        params["img_range"] = 255.0
         params["img_size"] = (126, 126)
         params["in_chans"] = 1
         params["upsampler"] = ""
         params["window_size"] = 7
     elif "color_jpeg" in name:
+        params["img_range"] = 255.0
         params["img_size"] = (126, 126)
         params["upsampler"] = ""
         params["window_size"] = 7