From 58c2fa4b2347177f3e8d050d9d879526317e4640 Mon Sep 17 00:00:00 2001
From: Tohrusky <65994850+Tohrusky@users.noreply.github.com>
Date: Fri, 28 Jun 2024 22:13:30 +0800
Subject: [PATCH] feat: add new model APISR-RRDB

---
 README.md                                     |  4 ++++
 src/Final2x_core/download_models.py           |  4 ++--
 src/Final2x_core/src/SRFactory/REALESRGAN.py  |  9 +++++++++
 src/Final2x_core/src/SRFactory/SRFactory.py   |  2 +-
 src/Final2x_core/src/SRncnn/REALESRGANncnn.py |  1 +
 tests/test_realesrgan.py                      | 13 +++++++++++++
 6 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 86a4e3b..1c39d53 100644
--- a/README.md
+++ b/README.md
@@ -94,6 +94,10 @@ Pass the config json string to the program through the `-j` parameter.
     - scale: 2, 3
     - noise: -1, 0, 3
 
+- APISR-RRDB:
+    - model: "APISR-RRDB"
+    - scale: 4
+
 - RealESRGAN-animevideov3:
     - model: "RealESRGAN-animevideov3"
     - scale: 2, 3, 4
diff --git a/src/Final2x_core/download_models.py b/src/Final2x_core/download_models.py
index 7f2a65a..c39615e 100644
--- a/src/Final2x_core/download_models.py
+++ b/src/Final2x_core/download_models.py
@@ -25,8 +25,8 @@
         "3392e740cda8c9e5308f33749da7ea76c1dbb972f4959db7745cfadbb8e0079f",
     ],
     "RealESRGAN": [
-        "https://github.com/Tohrusky/model-zoo/releases/download/v2.0.0/RealESRGAN.zip",
-        "79fff8b46b84c3f4b4da46837fb8f9630538e5936a8d9bc45697cceb6d29361f",
+        "https://github.com/Tohrusky/model-zoo/releases/download/v2.0.0/RealESRGAN-240628.zip",
+        "424ed35216c7045f28a48c04cb9bda96bba5e0abe45419eaf8914005d98c18fb",
     ],
     "Waifu2x": [
         "https://github.com/Tohrusky/model-zoo/releases/download/v2.0.0/Waifu2x.zip",
diff --git a/src/Final2x_core/src/SRFactory/REALESRGAN.py b/src/Final2x_core/src/SRFactory/REALESRGAN.py
index a0cf69d..034cbda 100644
--- a/src/Final2x_core/src/SRFactory/REALESRGAN.py
+++ b/src/Final2x_core/src/SRFactory/REALESRGAN.py
@@ -18,6 +18,7 @@ def _init_SR_class(self) -> None:
         #     2: {"param": "realesr-animevideov3-x4.param", "bin": "realesr-animevideov3-x4.bin", "scale": 4},
         #     3: {"param": "realesrgan-x4plus-anime.param", "bin": "realesrgan-x4plus-anime.bin", "scale": 4},
         #     4: {"param": "realesrgan-x4plus.param", "bin": "realesrgan-x4plus.bin", "scale": 4}
+        #     5: {"param": "APISR-RRDB-x4.param", "bin": "APISR-RRDB-x4.bin", "scale": 4}
         # }
         model_i = 0
         if self._model == "RealESRGAN-animevideov3":
@@ -48,6 +49,14 @@ def _init_SR_class(self) -> None:
                 model_i = 4
                 self._reset_modelscale(4)
 
+        elif self._model == "APISR-RRDB":
+            if self._modelscale == 4:
+                model_i = 5
+            else:
+                logger.warning("APISR-RRDB only support scale 4. Auto set to 4")
+                model_i = 5
+                self._reset_modelscale(4)
+
         else:
             logger.error("RealESRGAN model not implemented")
             raise NotImplementedError("RealESRGAN model not implemented")
diff --git a/src/Final2x_core/src/SRFactory/SRFactory.py b/src/Final2x_core/src/SRFactory/SRFactory.py
index 7bf49c9..9f39315 100644
--- a/src/Final2x_core/src/SRFactory/SRFactory.py
+++ b/src/Final2x_core/src/SRFactory/SRFactory.py
@@ -20,7 +20,7 @@ def getSR() -> Any:
 
             return REALCUGAN()
 
-        elif model in ["RealESRGAN-animevideov3", "RealESRGAN", "RealESRGAN-anime"]:
+        elif model in ["RealESRGAN-animevideov3", "RealESRGAN", "RealESRGAN-anime", "APISR-RRDB"]:
             from . import REALESRGAN
 
             return REALESRGAN()
diff --git a/src/Final2x_core/src/SRncnn/REALESRGANncnn.py b/src/Final2x_core/src/SRncnn/REALESRGANncnn.py
index e8c49e7..52243f5 100644
--- a/src/Final2x_core/src/SRncnn/REALESRGANncnn.py
+++ b/src/Final2x_core/src/SRncnn/REALESRGANncnn.py
@@ -34,6 +34,7 @@ def _load(
             2: {"param": "realesr-animevideov3-x4.param", "bin": "realesr-animevideov3-x4.bin", "scale": 4},
             3: {"param": "realesrgan-x4plus-anime.param", "bin": "realesrgan-x4plus-anime.bin", "scale": 4},
             4: {"param": "realesrgan-x4plus.param", "bin": "realesrgan-x4plus.bin", "scale": 4},
+            5: {"param": "APISR-RRDB-x4.param", "bin": "APISR-RRDB-x4.bin", "scale": 4},
         }
 
         param_path = _path / model_dict[self._model]["param"]  # type: ignore
diff --git a/tests/test_realesrgan.py b/tests/test_realesrgan.py
index acc0d4d..63dd21a 100644
--- a/tests/test_realesrgan.py
+++ b/tests/test_realesrgan.py
@@ -46,6 +46,19 @@ def test_case_RealESRGAN(self) -> None:
             assert calculate_image_similarity(img1, img2)
             assert compare_image_size(img1, img2, config.targetscale)
 
+    def test_case_APISR_RRDB(self) -> None:
+        from Final2x_core.src.SRFactory import REALESRGAN
+
+        config = getSRCONFIG()
+        config.model = "APISR-RRDB"  # type: ignore
+        for i in range(3, 6):
+            config.modelscale = i  # type: ignore
+            SR = REALESRGAN()
+            img1 = load_image()
+            img2 = SR.process(img1)
+            assert calculate_image_similarity(img1, img2)
+            assert compare_image_size(img1, img2, config.targetscale)
+
     def test_case_invalid_model(self) -> None:
         from Final2x_core.src.SRFactory import REALESRGAN