Add support for RetinexFormer

README.md

@@ -145,6 +145,10 @@ Spandrel currently supports a limited amount of network architectures. If the ar
 
 - [MixDehazeNet](https://github.com/AmeryXiong/MixDehazeNet) | [Models](https://drive.google.com/drive/folders/1ep6W4H3vNxshYjq71Tb3MzxrXGgaiM6C?usp=drive_link)
 
+#### Low-light Enhancement
+
+- [RetinexFormer](https://github.com/caiyuanhao1998/Retinexformer) | [Models](https://drive.google.com/drive/folders/1ynK5hfQachzc8y96ZumhkPPDXzHJwaQV?usp=drive_link)
+
 (All architectures marked with a `+` are only part of `spandrel_extra_arches`.)
 
 ## Security

libs/spandrel/spandrel/__helpers/main_registry.py

@@ -30,6 +30,7 @@
     OmniSR,
     RealCUGAN,
     RestoreFormer,
+    RetinexFormer,
     SCUNet,
     SwiftSRGAN,
     Swin2SR,
@@ -80,4 +81,5 @@
     ArchSupport.from_architecture(DRCT.DRCTArch()),
     ArchSupport.from_architecture(ESRGAN.ESRGANArch()),
     ArchSupport.from_architecture(PLKSR.PLKSRArch()),
+    ArchSupport.from_architecture(RetinexFormer.RetinexFormerArch()),
 )

libs/spandrel/spandrel/architectures/RetinexFormer/__init__.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import torch
+from typing_extensions import override
+
+from spandrel.util import KeyCondition, get_seq_len
+
+from ...__helpers.model_descriptor import (
+    Architecture,
+    ImageModelDescriptor,
+    ModelTiling,
+    SizeRequirements,
+    StateDict,
+)
+from .arch.retinexformer_arch import RetinexFormer
+
+
+def _call_fn(model: RetinexFormer, t: torch.Tensor) -> torch.Tensor:
+    h, w = t.shape[-2:]
+
+    if h < 3000 and w < 3000:
+        return model(t)
+
+    # this uses interlacing to split the image into 2 smaller parts
+    restored = torch.zeros_like(t)
+    restored[:, :, :, 1::2] = model(t[:, :, :, 1::2])
+    restored[:, :, :, 0::2] = model(t[:, :, :, 0::2])
+    return restored
+
+
+class RetinexFormerArch(Architecture[RetinexFormer]):
+    def __init__(self) -> None:
+        super().__init__(
+            id="RetinexFormer",
+            detect=KeyCondition.has_all(
+                "body.0.estimator.conv1.weight",
+                "body.0.estimator.conv1.bias",
+                "body.0.estimator.depth_conv.weight",
+                "body.0.estimator.depth_conv.bias",
+                "body.0.estimator.conv2.weight",
+                "body.0.estimator.conv2.bias",
+                "body.0.denoiser.embedding.weight",
+                "body.0.denoiser.mapping.weight",
+                "body.0.denoiser.encoder_layers.0.0.blocks.0.0.rescale",
+                "body.0.denoiser.encoder_layers.0.0.blocks.0.0.to_q.weight",
+                "body.0.denoiser.encoder_layers.0.0.blocks.0.0.to_v.weight",
+                "body.0.denoiser.encoder_layers.0.0.blocks.0.0.to_k.weight",
+                "body.0.denoiser.encoder_layers.0.0.blocks.0.0.proj.weight",
+                "body.0.denoiser.encoder_layers.0.0.blocks.0.0.pos_emb.0.weight",
+                "body.0.denoiser.encoder_layers.0.0.blocks.0.1.fn.net.0.weight",
+                "body.0.denoiser.encoder_layers.0.0.blocks.0.1.norm.weight",
+                "body.0.denoiser.encoder_layers.0.1.weight",
+                "body.0.denoiser.encoder_layers.0.2.weight",
+                "body.0.denoiser.bottleneck.blocks.0.0.rescale",
+                "body.0.denoiser.decoder_layers.0.0.weight",
+                "body.0.denoiser.decoder_layers.0.1.weight",
+                "body.0.denoiser.decoder_layers.0.2.blocks.0.0.rescale",
+            ),
+        )
+
+    @override
+    def load(self, state_dict: StateDict) -> ImageModelDescriptor[RetinexFormer]:
+        in_channels = 3
+        out_channels = 3
+        n_feat = 40
+        stage = 3
+        num_blocks = [1, 1, 1]
+
+        stage = get_seq_len(state_dict, "body")
+
+        n_feat = state_dict["body.0.denoiser.embedding.weight"].shape[0]
+        in_channels = state_dict["body.0.denoiser.embedding.weight"].shape[1]
+        out_channels = state_dict["body.0.denoiser.mapping.weight"].shape[0]
+
+        num_blocks = [
+            get_seq_len(state_dict, "body.0.denoiser.encoder_layers.0.0.blocks"),
+            get_seq_len(state_dict, "body.0.denoiser.encoder_layers.1.0.blocks"),
+            get_seq_len(state_dict, "body.0.denoiser.bottleneck.blocks"),
+        ]
+
+        model = RetinexFormer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            n_feat=n_feat,
+            stage=stage,
+            num_blocks=num_blocks,
+        )
+
+        return ImageModelDescriptor(
+            model,
+            state_dict,
+            architecture=self,
+            purpose="Restoration",
+            tags=[
+                f"{n_feat}nf",
+                f"{stage}s",
+                f"{num_blocks[0]}x{num_blocks[1]}x{num_blocks[2]}b",
+            ],
+            supports_half=False,  # TODO: verify
+            supports_bfloat16=True,
+            scale=1,
+            input_channels=in_channels,
+            output_channels=out_channels,
+            size_requirements=SizeRequirements(multiple_of=8),
+            tiling=ModelTiling.DISCOURAGED,
+            call_fn=_call_fn,
+        )

libs/spandrel/spandrel/architectures/RetinexFormer/arch/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Yuanhao Cai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.