From c756133541da478a35a74cda416d114a8973cf8e Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Wed, 27 Dec 2023 10:55:01 +0200
Subject: Add experimental HAT model

---
 modules/hat_model.py | 42 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)
 create mode 100644 modules/hat_model.py

(limited to 'modules/hat_model.py')

diff --git a/modules/hat_model.py b/modules/hat_model.py
new file mode 100644
index 00000000..553e1941
--- /dev/null
+++ b/modules/hat_model.py
@@ -0,0 +1,42 @@
+import os
+import sys
+
+from modules import modelloader, devices
+from modules.shared import opts
+from modules.upscaler import Upscaler, UpscalerData
+from modules.upscaler_utils import upscale_with_model
+
+
+class UpscalerHAT(Upscaler):
+    def __init__(self, dirname):
+        self.name = "HAT"
+        self.scalers = []
+        self.user_path = dirname
+        super().__init__()
+        for file in self.find_models(ext_filter=[".pt", ".pth"]):
+            name = modelloader.friendly_name(file)
+            scale = 4  # TODO: scale might not be 4, but we can't know without loading the model
+            scaler_data = UpscalerData(name, file, upscaler=self, scale=scale)
+            self.scalers.append(scaler_data)
+
+    def do_upscale(self, img, selected_model):
+        try:
+            model = self.load_model(selected_model)
+        except Exception as e:
+            print(f"Unable to load HAT model {selected_model}: {e}", file=sys.stderr)
+            return img
+        model.to(devices.device_esrgan)  # TODO: should probably be device_hat
+        return upscale_with_model(
+            model,
+            img,
+            tile_size=opts.ESRGAN_tile,  # TODO: should probably be HAT_tile
+            tile_overlap=opts.ESRGAN_tile_overlap,  # TODO: should probably be HAT_tile_overlap
+        )
+
+    def load_model(self, path: str):
+        if not os.path.isfile(path):
+            raise FileNotFoundError(f"Model file {path} not found")
+        return modelloader.load_spandrel_model(
+            path,
+            device=devices.device_esrgan,  # TODO: should probably be device_hat
+        )
-- 
cgit v1.2.3


From 4ad0c0c0a805da4bac03cff86ea17c25a1291546 Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Sat, 30 Dec 2023 16:37:03 +0200
Subject: Verify architecture for loaded Spandrel models

---
 modules/hat_model.py | 1 +
 1 file changed, 1 insertion(+)

(limited to 'modules/hat_model.py')

diff --git a/modules/hat_model.py b/modules/hat_model.py
index 553e1941..7f2abb41 100644
--- a/modules/hat_model.py
+++ b/modules/hat_model.py
@@ -39,4 +39,5 @@ class UpscalerHAT(Upscaler):
         return modelloader.load_spandrel_model(
             path,
             device=devices.device_esrgan,  # TODO: should probably be device_hat
+            expected_architecture='HAT',
         )
-- 
cgit v1.2.3