Add support for FBCNN

README.md

@@ -109,6 +109,10 @@ This has only been tested with the models that are linked here, and any unoffici
 
 - [SCUNet](https://github.com/cszn/SCUNet) | [GAN Model](https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_gan.pth) | [PSNR Model](https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth)
 
+#### DeJPEG
+
+- [FBCNN](https://github.com/jiaxi-jiang/FBCNN) | [Models](https://github.com/jiaxi-jiang/FBCNN/releases/tag/v1.0)
+
 
 ## Contributing
 

src/spandrel/__helpers/main_registry.py

@@ -5,6 +5,7 @@
 from ..architectures import (
     DAT,
     ESRGAN,
+    FBCNN,
     GFPGAN,
     HAT,
     MAT,
@@ -153,6 +154,18 @@ def _detect(state_dict: StateDict) -> bool:
         detect=_has_keys("m_head.0.weight", "m_tail.0.weight"),
         load=SCUNet.load,
     ),
+    ArchSupport(
+        id="FBCNN",
+        detect=_has_keys(
+            "m_head.weight",
+            "m_down1.0.res.0.weight",
+            "m_down2.0.res.0.weight",
+            "m_body_encoder.0.res.0.weight",
+            "m_tail.weight",
+            "qf_pred.0.res.0.weight",
+        ),
+        load=FBCNN.load,
+    ),
     ArchSupport(
         id="DAT",
         detect=_has_keys("layers.0.blocks.2.attn.attn_mask_0", "conv_first.weight"),

src/spandrel/architectures/FBCNN/__init__.py

@@ -0,0 +1,51 @@
+from ...__helpers.model_descriptor import (
+    RestorationModelDescriptor,
+    SizeRequirements,
+    StateDict,
+)
+from .arch.FBCNN import FBCNN
+
+
+def load(state_dict: StateDict) -> RestorationModelDescriptor[FBCNN]:
+    in_nc = 3
+    out_nc = 3
+    nc = [64, 128, 256, 512]
+    nb = 4
+    act_mode = "R"
+    downsample_mode = "strideconv"
+    upsample_mode = "convtranspose"
+
+    in_nc = state_dict["m_head.weight"].shape[1]
+    out_nc = state_dict["m_tail.weight"].shape[0]
+
+    for i in range(0, 20):
+        if f"m_down1.{i}.weight" in state_dict:
+            nb = i
+            break
+
+    nc[0] = state_dict["m_head.weight"].shape[0]
+    nc[1] = state_dict[f"m_down1.{nb}.weight"].shape[0]
+    nc[2] = state_dict[f"m_down2.{nb}.weight"].shape[0]
+    nc[3] = state_dict[f"m_down3.{nb}.weight"].shape[0]
+
+    model = FBCNN(
+        in_nc=in_nc,
+        out_nc=out_nc,
+        nc=nc,
+        nb=nb,
+        act_mode=act_mode,
+        downsample_mode=downsample_mode,
+        upsample_mode=upsample_mode,
+    )
+
+    return RestorationModelDescriptor(
+        model,
+        state_dict,
+        architecture="FBCNN",
+        tags=[],
+        supports_half=True,  # TODO
+        supports_bfloat16=True,  # TODO
+        input_channels=in_nc,
+        output_channels=out_nc,
+        size=SizeRequirements(minimum=16),  # TODO
+    )