diff options
| author | Grafting Rayman <156515434+GraftingRayman@users.noreply.github.com> | 2025-01-17 11:00:30 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2025-01-17 11:00:30 +0000 |
| commit | 495ffc4777522e40941753e3b1b79c02f84b25b4 (patch) | |
| tree | 5130fcb8676afdcb619a5e5eaef3ac28e135bc08 /r_basicsr/models/swinir_model.py | |
| parent | febd45814cd41560c5247aacb111d8d013f3a303 (diff) | |
| download | Comfyui-reactor-node-495ffc4777522e40941753e3b1b79c02f84b25b4.tar.gz | |
Add files via upload
Diffstat (limited to 'r_basicsr/models/swinir_model.py')
| -rw-r--r-- | r_basicsr/models/swinir_model.py | 33 |
1 files changed, 33 insertions, 0 deletions
diff --git a/r_basicsr/models/swinir_model.py b/r_basicsr/models/swinir_model.py new file mode 100644 index 0000000..18e5550 --- /dev/null +++ b/r_basicsr/models/swinir_model.py @@ -0,0 +1,33 @@ +import torch
+from torch.nn import functional as F
+
+from r_basicsr.utils.registry import MODEL_REGISTRY
+from .sr_model import SRModel
+
+
+@MODEL_REGISTRY.register()
+class SwinIRModel(SRModel):
+
+ def test(self):
+ # pad to multiplication of window_size
+ window_size = self.opt['network_g']['window_size']
+ scale = self.opt.get('scale', 1)
+ mod_pad_h, mod_pad_w = 0, 0
+ _, _, h, w = self.lq.size()
+ if h % window_size != 0:
+ mod_pad_h = window_size - h % window_size
+ if w % window_size != 0:
+ mod_pad_w = window_size - w % window_size
+ img = F.pad(self.lq, (0, mod_pad_w, 0, mod_pad_h), 'reflect')
+ if hasattr(self, 'net_g_ema'):
+ self.net_g_ema.eval()
+ with torch.no_grad():
+ self.output = self.net_g_ema(img)
+ else:
+ self.net_g.eval()
+ with torch.no_grad():
+ self.output = self.net_g(img)
+ self.net_g.train()
+
+ _, _, h, w = self.output.size()
+ self.output = self.output[:, :, 0:h - mod_pad_h * scale, 0:w - mod_pad_w * scale]
|