From 55c3029c51d6eb95a36f5c2f2865aceb64fa4321 Mon Sep 17 00:00:00 2001
From: hyw-dev <1315176066@qq.com>
Date: Fri, 8 Nov 2024 17:59:45 +0800
Subject: [PATCH 1/5] implement DAT.

---
 ccrestoration/arch/__init__.py     |    1 +
 ccrestoration/arch/dat_arch.py     | 1049 ++++++++++++++++++++++++++++
 ccrestoration/config/__init__.py   |    1 +
 ccrestoration/config/dat_config.py |  172 +++++
 ccrestoration/model/__init__.py    |    1 +
 ccrestoration/model/dat_model.py   |   45 ++
 ccrestoration/type/arch.py         |    1 +
 ccrestoration/type/config.py       |   14 +
 ccrestoration/type/model.py        |    1 +
 tests/test_dat.py                  |   29 +
 10 files changed, 1314 insertions(+)
 create mode 100644 ccrestoration/arch/dat_arch.py
 create mode 100644 ccrestoration/config/dat_config.py
 create mode 100644 ccrestoration/model/dat_model.py
 create mode 100644 tests/test_dat.py

diff --git a/ccrestoration/arch/__init__.py b/ccrestoration/arch/__init__.py
index 908ab3e..596d872 100644
--- a/ccrestoration/arch/__init__.py
+++ b/ccrestoration/arch/__init__.py
@@ -13,3 +13,4 @@
 from ccrestoration.arch.iconvsr_arch import IconVSR  # noqa
 from ccrestoration.arch.msrswvsr_arch import MSRSWVSR  # noqa
 from ccrestoration.arch.scunet_arch import SCUNet  # noqa
+from ccrestoration.arch.dat_arch import DAT  # noqa
diff --git a/ccrestoration/arch/dat_arch.py b/ccrestoration/arch/dat_arch.py
new file mode 100644
index 0000000..d7f8d87
--- /dev/null
+++ b/ccrestoration/arch/dat_arch.py
@@ -0,0 +1,1049 @@
+# type: ignore  # noqa: PGH003
+import numpy as np
+import torch
+from einops import rearrange
+from einops.layers.torch import Rearrange
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.init import trunc_normal_
+from torch.utils import checkpoint
+
+from ccrestoration.arch import ARCH_REGISTRY
+from ccrestoration.arch.arch_util import Upsample, DropPath
+from ccrestoration.type import ArchType
+
+
+@ARCH_REGISTRY.register(name=ArchType.DAT)
+class DAT(nn.Module):
+    """Dual Aggregation Transformer
+    Args:
+        img_size (int): Input image size. Default: 64
+        in_chans (int): Number of input image channels. Default: 3
+        embed_dim (int): Patch embedding dimension. Default: 180
+        depths (tuple(int)): Depth of each residual group (number of DATB in each RG).
+        split_size (tuple(int)): Height and Width of spatial window.
+        num_heads (tuple(int)): Number of attention heads in different residual groups.
+        expansion_factor (float): Ratio of ffn hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        act_layer (nn.Module): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm
+        use_chk (bool): Whether to use checkpointing to save memory.
+        upscale: Upscale factor. 2/3/4 for image SR
+        img_range: Image range. 1. or 255.
+        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
+    """
+
+    def __init__(
+        self,
+        img_size=64,
+        in_chans=3,
+        embed_dim=180,
+        split_size=[2, 4],
+        depth=[2, 2, 2, 2],
+        num_heads=[2, 2, 2, 2],
+        expansion_factor=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.1,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        use_chk=False,
+        upscale=2,
+        img_range=1.0,
+        resi_connection="1conv",
+        upsampler="pixelshuffle",
+        **kwargs,
+    ):
+        super().__init__()
+
+        num_in_ch = in_chans
+        num_out_ch = in_chans
+        num_feat = 64
+        self.img_range = img_range
+        if in_chans == 3:
+            rgb_mean = (0.5, 0.5, 0.5)
+            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
+        else:
+            self.mean = torch.zeros(1, 1, 1, 1)
+        self.upscale = upscale
+        self.upsampler = upsampler
+
+        # ------------------------- 1, Shallow Feature Extraction ------------------------- #
+        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)
+
+        # ------------------------- 2, Deep Feature Extraction ------------------------- #
+        self.num_layers = len(depth)
+        self.use_chk = use_chk
+        self.num_features = self.embed_dim = (
+            embed_dim  # num_features for consistency with other models
+        )
+        heads = num_heads
+
+        self.before_RG = nn.Sequential(
+            Rearrange("b c h w -> b (h w) c"), nn.LayerNorm(embed_dim)
+        )
+
+        curr_dim = embed_dim
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, np.sum(depth))
+        ]  # stochastic depth decay rule
+
+        self.layers = nn.ModuleList()
+        for i in range(self.num_layers):
+            layer = ResidualGroup(
+                dim=embed_dim,
+                num_heads=heads[i],
+                reso=img_size,
+                split_size=split_size,
+                expansion_factor=expansion_factor,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_paths=dpr[sum(depth[:i]) : sum(depth[: i + 1])],
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                depth=depth[i],
+                use_chk=use_chk,
+                resi_connection=resi_connection,
+                rg_idx=i,
+            )
+            self.layers.append(layer)
+
+        self.norm = norm_layer(curr_dim)
+        # build the last conv layer in deep feature extraction
+        if resi_connection == "1conv":
+            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
+        elif resi_connection == "3conv":
+            # to save parameters and memory
+            self.conv_after_body = nn.Sequential(
+                nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1),
+            )
+
+        # ------------------------- 3, Reconstruction ------------------------- #
+        if self.upsampler == "pixelshuffle":
+            # for classical SR
+            self.conv_before_upsample = nn.Sequential(
+                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
+            )
+            self.upsample = Upsample(upscale, num_feat)
+            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+        elif self.upsampler == "pixelshuffledirect":
+            # for lightweight SR (to save parameters)
+            self.upsample = UpsampleOneStep(
+                upscale, embed_dim, num_out_ch, (img_size, img_size)
+            )
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(
+            m, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d)
+        ):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward_features(self, x):
+        _, _, H, W = x.shape
+        x_size = [H, W]
+        x = self.before_RG(x)
+        for layer in self.layers:
+            x = layer(x, x_size)
+        x = self.norm(x)
+        x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W)
+
+        return x
+
+    def forward(self, x):
+        """Input: x: (B, C, H, W)"""
+        self.mean = self.mean.type_as(x)
+        x = (x - self.mean) * self.img_range
+
+        if self.upsampler == "pixelshuffle":
+            # for image SR
+            x = self.conv_first(x)
+            x = self.conv_after_body(self.forward_features(x)) + x
+            x = self.conv_before_upsample(x)
+            x = self.conv_last(self.upsample(x))
+        elif self.upsampler == "pixelshuffledirect":
+            # for lightweight SR
+            x = self.conv_first(x)
+            x = self.conv_after_body(self.forward_features(x)) + x
+            x = self.upsample(x)
+
+        x = x / self.img_range + self.mean
+        return x
+
+def img2windows(img, H_sp, W_sp):
+    """Input: Image (B, C, H, W)
+    Output: Window Partition (B', N, C)
+    """
+    B, C, H, W = img.shape
+    img_reshape = img.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp)
+    img_perm = (
+        img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, H_sp * W_sp, C)
+    )
+    return img_perm
+
+
+def windows2img(img_splits_hw, H_sp, W_sp, H, W):
+    """Input: Window Partition (B', N, C)
+    Output: Image (B, H, W, C)
+    """
+    B = int(img_splits_hw.shape[0] / (H * W / H_sp / W_sp))
+
+    img = img_splits_hw.view(B, H // H_sp, W // W_sp, H_sp, W_sp, -1)
+    img = img.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return img
+
+
+class SpatialGate(nn.Module):
+    """Spatial-Gate.
+
+    Args:
+    ----
+        dim (int): Half of input channels.
+
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.conv = nn.Conv2d(
+            dim, dim, kernel_size=3, stride=1, padding=1, groups=dim
+        )  # DW Conv
+
+    def forward(self, x, H, W):
+        # Split
+        x1, x2 = x.chunk(2, dim=-1)
+        B, N, C = x.shape
+        x2 = (
+            self.conv(self.norm(x2).transpose(1, 2).contiguous().view(B, C // 2, H, W))
+            .flatten(2)
+            .transpose(-1, -2)
+            .contiguous()
+        )
+
+        return x1 * x2
+
+
+class SGFN(nn.Module):
+    """Spatial-Gate Feed-Forward Network.
+
+    Args:
+    ----
+        in_features (int): Number of input channels.
+        hidden_features (int | None): Number of hidden channels. Default: None
+        out_features (int | None): Number of output channels. Default: None
+        act_layer (nn.Module): Activation layer. Default: nn.GELU
+        drop (float): Dropout rate. Default: 0.0
+
+    """
+
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.sg = SpatialGate(hidden_features // 2)
+        self.fc2 = nn.Linear(hidden_features // 2, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x, H, W):
+        """Input: x: (B, H*W, C), H, W
+        Output: x: (B, H*W, C)
+        """
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+
+        x = self.sg(x, H, W)
+        x = self.drop(x)
+
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class DynamicPosBias(nn.Module):
+    # The implementation builds on Crossformer code https://github.com/cheerss/CrossFormer/blob/main/models/crossformer.py
+
+    """Dynamic Relative Position Bias.
+
+    Args:
+    ----
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        residual (bool):  If True, use residual strage to connect conv.
+
+    """
+
+    def __init__(self, dim, num_heads, residual):
+        super().__init__()
+        self.residual = residual
+        self.num_heads = num_heads
+        self.pos_dim = dim // 4
+        self.pos_proj = nn.Linear(2, self.pos_dim)
+        self.pos1 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.pos_dim),
+        )
+        self.pos2 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.pos_dim),
+        )
+        self.pos3 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.num_heads),
+        )
+
+    def forward(self, biases):
+        if self.residual:
+            pos = self.pos_proj(biases)  # 2Gh-1 * 2Gw-1, heads
+            pos = pos + self.pos1(pos)
+            pos = pos + self.pos2(pos)
+            pos = self.pos3(pos)
+        else:
+            pos = self.pos3(self.pos2(self.pos1(self.pos_proj(biases))))
+        return pos
+
+
+class Spatial_Attention(nn.Module):
+    """Spatial Window Self-Attention.
+    It supports rectangle window (containing square window).
+
+    Args:
+    ----
+        dim (int): Number of input channels.
+        idx (int): The indentix of window. (0/1)
+        split_size (tuple(int)): Height and Width of spatial window.
+        dim_out (int | None): The dimension of the attention output. Default: None
+        num_heads (int): Number of attention heads. Default: 6
+        attn_drop (float): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float): Dropout ratio of output. Default: 0.0
+        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set
+        position_bias (bool): The dynamic relative position bias. Default: True
+
+    """
+
+    def __init__(
+        self,
+        dim,
+        idx,
+        split_size=[8, 8],
+        dim_out=None,
+        num_heads=6,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        qk_scale=None,
+        position_bias=True,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.dim_out = dim_out or dim
+        self.split_size = split_size
+        self.num_heads = num_heads
+        self.idx = idx
+        self.position_bias = position_bias
+
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        if idx == 0:
+            H_sp, W_sp = self.split_size[0], self.split_size[1]
+        elif idx == 1:
+            W_sp, H_sp = self.split_size[0], self.split_size[1]
+        else:
+            print("ERROR MODE", idx)
+            exit(0)
+        self.H_sp = H_sp
+        self.W_sp = W_sp
+
+        if self.position_bias:
+            self.pos = DynamicPosBias(self.dim // 4, self.num_heads, residual=False)
+            # generate mother-set
+            position_bias_h = torch.arange(1 - self.H_sp, self.H_sp)
+            position_bias_w = torch.arange(1 - self.W_sp, self.W_sp)
+            biases = torch.stack(
+                torch.meshgrid([position_bias_h, position_bias_w], indexing="ij")
+            )
+            biases = biases.flatten(1).transpose(0, 1).contiguous().float()
+            self.register_buffer("rpe_biases", biases)
+
+            # get pair-wise relative position index for each token inside the window
+            coords_h = torch.arange(self.H_sp)
+            coords_w = torch.arange(self.W_sp)
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))
+            coords_flatten = torch.flatten(coords, 1)
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+            relative_coords[:, :, 0] += self.H_sp - 1
+            relative_coords[:, :, 1] += self.W_sp - 1
+            relative_coords[:, :, 0] *= 2 * self.W_sp - 1
+            relative_position_index = relative_coords.sum(-1)
+            self.register_buffer("relative_position_index", relative_position_index)
+
+        self.attn_drop = nn.Dropout(attn_drop)
+
+    def im2win(self, x, H, W):
+        B, N, C = x.shape
+        x = x.transpose(-2, -1).contiguous().view(B, C, H, W)
+        x = img2windows(x, self.H_sp, self.W_sp)
+        x = (
+            x.reshape(-1, self.H_sp * self.W_sp, self.num_heads, C // self.num_heads)
+            .permute(0, 2, 1, 3)
+            .contiguous()
+        )
+        return x
+
+    def forward(self, qkv, H, W, mask=None):
+        """Input: qkv: (B, 3*L, C), H, W, mask: (B, N, N), N is the window size
+        Output: x (B, H, W, C)
+        """
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        B, L, C = q.shape
+        assert L == H * W, "flatten img_tokens has wrong size"
+
+        # partition the q,k,v, image to window
+        q = self.im2win(q, H, W)
+        k = self.im2win(k, H, W)
+        v = self.im2win(v, H, W)
+
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)  # B head N C @ B head C N --> B head N N
+
+        # calculate drpe
+        if self.position_bias:
+            pos = self.pos(self.rpe_biases)
+            # select position bias
+            relative_position_bias = pos[self.relative_position_index.view(-1)].view(
+                self.H_sp * self.W_sp, self.H_sp * self.W_sp, -1
+            )
+            relative_position_bias = relative_position_bias.permute(
+                2, 0, 1
+            ).contiguous()
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+        N = attn.shape[3]
+
+        # use mask for shift window
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(
+                0
+            )
+            attn = attn.view(-1, self.num_heads, N, N)
+
+        attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype)
+        attn = self.attn_drop(attn)
+
+        x = attn @ v
+        x = x.transpose(1, 2).reshape(
+            -1, self.H_sp * self.W_sp, C
+        )  # B head N N @ B head N C
+
+        # merge the window, window to image
+        x = windows2img(x, self.H_sp, self.W_sp, H, W)  # B H' W' C
+
+        return x
+
+
+class Axial_Spatial_Attention(nn.Module):
+    """Axial Spatial Self-Attention
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads. Default: 6
+        split_size (tuple(int)): Height and Width of spatial window.
+        shift_size (tuple(int)): Shift size for spatial window.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float): Dropout rate. Default: 0.0
+        attn_drop (float): Attention dropout rate. Default: 0.0
+        rg_idx (int): The indentix of Residual Group (RG)
+        b_idx (int): The indentix of Block in each RG
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        reso=64,
+        split_size=[8, 8],
+        shift_size=[1, 2],
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        rg_idx=0,
+        b_idx=0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.split_size = split_size
+        self.shift_size = shift_size
+        self.b_idx = b_idx
+        self.rg_idx = rg_idx
+        self.patches_resolution = reso
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+
+        assert (
+            0 <= self.shift_size[0] < self.split_size[0]
+        ), "shift_size must in 0-split_size0"
+        assert (
+            0 <= self.shift_size[1] < self.split_size[1]
+        ), "shift_size must in 0-split_size1"
+
+        self.branch_num = 2
+
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(drop)
+
+        self.attns = nn.ModuleList([
+            Spatial_Attention(
+                dim // 2,
+                idx=i,
+                split_size=split_size,
+                num_heads=num_heads // 2,
+                dim_out=dim // 2,
+                qk_scale=qk_scale,
+                attn_drop=attn_drop,
+                proj_drop=drop,
+                position_bias=True,
+            )
+            for i in range(self.branch_num)
+        ])
+
+        if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or (
+            self.rg_idx % 2 != 0 and self.b_idx % 4 == 0
+        ):
+            attn_mask = self.calculate_mask(
+                self.patches_resolution, self.patches_resolution
+            )
+            self.register_buffer("attn_mask_0", attn_mask[0])
+            self.register_buffer("attn_mask_1", attn_mask[1])
+        else:
+            attn_mask = None
+            self.register_buffer("attn_mask_0", None)
+            self.register_buffer("attn_mask_1", None)
+
+        self.dwconv = nn.Sequential(
+            nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim),
+            nn.BatchNorm2d(dim),
+            nn.GELU(),
+        )
+        self.channel_interaction = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(dim, dim // 8, kernel_size=1),
+            nn.BatchNorm2d(dim // 8),
+            nn.GELU(),
+            nn.Conv2d(dim // 8, dim, kernel_size=1),
+        )
+        self.spatial_interaction = nn.Sequential(
+            nn.Conv2d(dim, dim // 16, kernel_size=1),
+            nn.BatchNorm2d(dim // 16),
+            nn.GELU(),
+            nn.Conv2d(dim // 16, 1, kernel_size=1),
+        )
+
+    def calculate_mask(self, H, W):
+        # The implementation builds on Swin Transformer code https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
+        # calculate attention mask for shift window
+        img_mask_0 = torch.zeros((1, H, W, 1))  # 1 H W 1 idx=0
+        img_mask_1 = torch.zeros((1, H, W, 1))  # 1 H W 1 idx=1
+        h_slices_0 = (
+            slice(0, -self.split_size[0]),
+            slice(-self.split_size[0], -self.shift_size[0]),
+            slice(-self.shift_size[0], None),
+        )
+        w_slices_0 = (
+            slice(0, -self.split_size[1]),
+            slice(-self.split_size[1], -self.shift_size[1]),
+            slice(-self.shift_size[1], None),
+        )
+
+        h_slices_1 = (
+            slice(0, -self.split_size[1]),
+            slice(-self.split_size[1], -self.shift_size[1]),
+            slice(-self.shift_size[1], None),
+        )
+        w_slices_1 = (
+            slice(0, -self.split_size[0]),
+            slice(-self.split_size[0], -self.shift_size[0]),
+            slice(-self.shift_size[0], None),
+        )
+        cnt = 0
+        for h in h_slices_0:
+            for w in w_slices_0:
+                img_mask_0[:, h, w, :] = cnt
+                cnt += 1
+        cnt = 0
+        for h in h_slices_1:
+            for w in w_slices_1:
+                img_mask_1[:, h, w, :] = cnt
+                cnt += 1
+
+        # calculate mask for window-0
+        img_mask_0 = img_mask_0.view(
+            1,
+            H // self.split_size[0],
+            self.split_size[0],
+            W // self.split_size[1],
+            self.split_size[1],
+            1,
+        )
+        img_mask_0 = (
+            img_mask_0.permute(0, 1, 3, 2, 4, 5)
+            .contiguous()
+            .view(-1, self.split_size[0], self.split_size[1], 1)
+        )  # nW, sw[0], sw[1], 1
+        mask_windows_0 = img_mask_0.view(-1, self.split_size[0] * self.split_size[1])
+        attn_mask_0 = mask_windows_0.unsqueeze(1) - mask_windows_0.unsqueeze(2)
+        attn_mask_0 = attn_mask_0.masked_fill(attn_mask_0 != 0, -100.0).masked_fill(
+            attn_mask_0 == 0, 0.0
+        )
+
+        # calculate mask for window-1
+        img_mask_1 = img_mask_1.view(
+            1,
+            H // self.split_size[1],
+            self.split_size[1],
+            W // self.split_size[0],
+            self.split_size[0],
+            1,
+        )
+        img_mask_1 = (
+            img_mask_1.permute(0, 1, 3, 2, 4, 5)
+            .contiguous()
+            .view(-1, self.split_size[1], self.split_size[0], 1)
+        )  # nW, sw[1], sw[0], 1
+        mask_windows_1 = img_mask_1.view(-1, self.split_size[1] * self.split_size[0])
+        attn_mask_1 = mask_windows_1.unsqueeze(1) - mask_windows_1.unsqueeze(2)
+        attn_mask_1 = attn_mask_1.masked_fill(attn_mask_1 != 0, -100.0).masked_fill(
+            attn_mask_1 == 0, 0.0
+        )
+
+        return attn_mask_0, attn_mask_1
+
+    def forward(self, x, H, W):
+        """Input: x: (B, H*W, C), H, W
+        Output: x: (B, H*W, C)
+        """
+        B, L, C = x.shape
+        assert L == H * W, "flatten img_tokens has wrong size"
+
+        qkv = self.qkv(x).reshape(B, -1, 3, C).permute(2, 0, 1, 3)  # 3, B, HW, C
+        # V without partition
+        v = qkv[2].transpose(-2, -1).contiguous().view(B, C, H, W)
+
+        # image padding
+        max_split_size = max(self.split_size[0], self.split_size[1])
+        pad_l = pad_t = 0
+        pad_r = (max_split_size - W % max_split_size) % max_split_size
+        pad_b = (max_split_size - H % max_split_size) % max_split_size
+
+        qkv = qkv.reshape(3 * B, H, W, C).permute(0, 3, 1, 2)  # 3B C H W
+        qkv = (
+            F.pad(qkv, (pad_l, pad_r, pad_t, pad_b))
+            .reshape(3, B, C, -1)
+            .transpose(-2, -1)
+        )  # l r t b
+        _H = pad_b + H
+        _W = pad_r + W
+        _L = _H * _W
+
+        # window-0 and window-1 on split channels [C/2, C/2]; for square windows (e.g., 8x8), window-0 and window-1 can be merged
+        # shift in block: (0, 4, 8, ...), (2, 6, 10, ...), (0, 4, 8, ...), (2, 6, 10, ...), ...
+        if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or (
+            self.rg_idx % 2 != 0 and self.b_idx % 4 == 0
+        ):
+            qkv = qkv.view(3, B, _H, _W, C)
+            qkv_0 = torch.roll(
+                qkv[:, :, :, :, : C // 2],
+                shifts=(-self.shift_size[0], -self.shift_size[1]),
+                dims=(2, 3),
+            )
+            qkv_0 = qkv_0.view(3, B, _L, C // 2)
+            qkv_1 = torch.roll(
+                qkv[:, :, :, :, C // 2 :],
+                shifts=(-self.shift_size[1], -self.shift_size[0]),
+                dims=(2, 3),
+            )
+            qkv_1 = qkv_1.view(3, B, _L, C // 2)
+
+            if self.patches_resolution != _H or self.patches_resolution != _W:
+                mask_tmp = self.calculate_mask(_H, _W)
+                x1_shift = self.attns[0](qkv_0, _H, _W, mask=mask_tmp[0].to(x.device))
+                x2_shift = self.attns[1](qkv_1, _H, _W, mask=mask_tmp[1].to(x.device))
+            else:
+                x1_shift = self.attns[0](qkv_0, _H, _W, mask=self.attn_mask_0)
+                x2_shift = self.attns[1](qkv_1, _H, _W, mask=self.attn_mask_1)
+
+            x1 = torch.roll(
+                x1_shift, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2)
+            )
+            x2 = torch.roll(
+                x2_shift, shifts=(self.shift_size[1], self.shift_size[0]), dims=(1, 2)
+            )
+            x1 = x1[:, :H, :W, :].reshape(B, L, C // 2)
+            x2 = x2[:, :H, :W, :].reshape(B, L, C // 2)
+            # attention output
+            attened_x = torch.cat([x1, x2], dim=2)
+
+        else:
+            x1 = self.attns[0](qkv[:, :, :, : C // 2], _H, _W)[:, :H, :W, :].reshape(
+                B, L, C // 2
+            )
+            x2 = self.attns[1](qkv[:, :, :, C // 2 :], _H, _W)[:, :H, :W, :].reshape(
+                B, L, C // 2
+            )
+            # attention output
+            attened_x = torch.cat([x1, x2], dim=2)
+
+        # convolution output
+        conv_x = self.dwconv(v)
+
+        # Adaptive Interaction Module (AIM)
+        # C-Map (before sigmoid)
+        channel_map = (
+            self.channel_interaction(conv_x)
+            .permute(0, 2, 3, 1)
+            .contiguous()
+            .view(B, 1, C)
+        )
+        # S-Map (before sigmoid)
+        attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W)
+        spatial_map = self.spatial_interaction(attention_reshape)
+
+        # C-I
+        attened_x = attened_x * torch.sigmoid(channel_map)
+        # S-I
+        conv_x = torch.sigmoid(spatial_map) * conv_x
+        conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(B, L, C)
+
+        x = attened_x + conv_x
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class Axial_Channel_Attention(nn.Module):
+    # The implementation builds on XCiT code https://github.com/facebookresearch/xcit
+
+    """Axial Channel Self-Attention
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads. Default: 6
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set.
+        attn_drop (float): Attention dropout rate. Default: 0.0
+        drop_path (float): Stochastic depth rate. Default: 0.0
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.dwconv = nn.Sequential(
+            nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim),
+            nn.BatchNorm2d(dim),
+            nn.GELU(),
+        )
+        self.channel_interaction = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(dim, dim // 8, kernel_size=1),
+            nn.BatchNorm2d(dim // 8),
+            nn.GELU(),
+            nn.Conv2d(dim // 8, dim, kernel_size=1),
+        )
+        self.spatial_interaction = nn.Sequential(
+            nn.Conv2d(dim, dim // 16, kernel_size=1),
+            nn.BatchNorm2d(dim // 16),
+            nn.GELU(),
+            nn.Conv2d(dim // 16, 1, kernel_size=1),
+        )
+
+    def forward(self, x, H, W):
+        """Input: x: (B, H*W, C), H, W
+        Output: x: (B, H*W, C)
+        """
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        q = q.transpose(-2, -1)
+        k = k.transpose(-2, -1)
+        v = v.transpose(-2, -1)
+
+        v_ = v.reshape(B, C, N).contiguous().view(B, C, H, W)
+
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        # attention output
+        attened_x = (attn @ v).permute(0, 3, 1, 2).reshape(B, N, C)
+
+        # convolution output
+        conv_x = self.dwconv(v_)
+
+        # Adaptive Interaction Module (AIM)
+        # C-Map (before sigmoid)
+        attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W)
+        channel_map = self.channel_interaction(attention_reshape)
+        # S-Map (before sigmoid)
+        spatial_map = (
+            self.spatial_interaction(conv_x)
+            .permute(0, 2, 3, 1)
+            .contiguous()
+            .view(B, N, 1)
+        )
+
+        # S-I
+        attened_x = attened_x * torch.sigmoid(spatial_map)
+        # C-I
+        conv_x = conv_x * torch.sigmoid(channel_map)
+        conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(B, N, C)
+
+        x = attened_x + conv_x
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class DATB(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        reso=64,
+        split_size=[2, 4],
+        shift_size=[1, 2],
+        expansion_factor=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        rg_idx=0,
+        b_idx=0,
+    ):
+        super().__init__()
+
+        self.norm1 = norm_layer(dim)
+
+        if b_idx % 2 == 0:
+            # DSTB
+            self.attn = Axial_Spatial_Attention(
+                dim,
+                num_heads=num_heads,
+                reso=reso,
+                split_size=split_size,
+                shift_size=shift_size,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop,
+                attn_drop=attn_drop,
+                rg_idx=rg_idx,
+                b_idx=b_idx,
+            )
+        else:
+            # DCTB
+            self.attn = Axial_Channel_Attention(
+                dim,
+                num_heads=num_heads,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                attn_drop=attn_drop,
+                proj_drop=drop,
+            )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        ffn_hidden_dim = int(dim * expansion_factor)
+        self.ffn = SGFN(
+            in_features=dim,
+            hidden_features=ffn_hidden_dim,
+            out_features=dim,
+            act_layer=act_layer,
+        )
+        self.norm2 = norm_layer(dim)
+
+    def forward(self, x, x_size):
+        """Input: x: (B, H*W, C), x_size: (H, W)
+        Output: x: (B, H*W, C)
+        """
+        H, W = x_size
+        x = x + self.drop_path(self.attn(self.norm1(x), H, W))
+        x = x + self.drop_path(self.ffn(self.norm2(x), H, W))
+
+        return x
+
+
+class ResidualGroup(nn.Module):
+    """ResidualGroup
+    Args:
+        dim (int): Number of input channels.
+        reso (int): Input resolution.
+        num_heads (int): Number of attention heads.
+        split_size (tuple(int)): Height and Width of spatial window.
+        expansion_factor (float): Ratio of ffn hidden dim to embedding dim.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop (float): Dropout rate. Default: 0
+        attn_drop(float): Attention dropout rate. Default: 0
+        drop_paths (float | None): Stochastic depth rate.
+        act_layer (nn.Module): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm
+        depth (int): Number of dual aggregation Transformer blocks in residual group.
+        use_chk (bool): Whether to use checkpointing to save memory.
+        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
+    """
+
+    def __init__(
+        self,
+        dim,
+        reso,
+        num_heads,
+        split_size=[2, 4],
+        expansion_factor=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_paths=None,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        depth=2,
+        use_chk=False,
+        resi_connection="1conv",
+        rg_idx=0,
+    ):
+        super().__init__()
+        self.use_chk = use_chk
+        self.reso = reso
+
+        self.blocks = nn.ModuleList([
+            DATB(
+                dim=dim,
+                num_heads=num_heads,
+                reso=reso,
+                split_size=split_size,
+                shift_size=[split_size[0] // 2, split_size[1] // 2],
+                expansion_factor=expansion_factor,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop,
+                attn_drop=attn_drop,
+                drop_path=drop_paths[i],
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                rg_idx=rg_idx,
+                b_idx=i,
+            )
+            for i in range(depth)
+        ])
+
+        if resi_connection == "1conv":
+            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
+        elif resi_connection == "3conv":
+            self.conv = nn.Sequential(
+                nn.Conv2d(dim, dim // 4, 3, 1, 1),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(dim // 4, dim, 3, 1, 1),
+            )
+
+    def forward(self, x, x_size):
+        """Input: x: (B, H*W, C), x_size: (H, W)
+        Output: x: (B, H*W, C)
+        """
+        H, W = x_size
+        res = x
+        for blk in self.blocks:
+            if self.use_chk:
+                x = checkpoint.checkpoint(blk, x, x_size, use_reentrant=False)
+            else:
+                x = blk(x, x_size)
+        x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W)
+        x = self.conv(x)
+        x = rearrange(x, "b c h w -> b (h w) c")
+        x = res + x
+
+        return x
+
+
+class UpsampleOneStep(nn.Sequential):
+    """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
+       Used in lightweight SR to save parameters.
+
+    Args:
+    ----
+        scale (int): Scale factor. Supported scales: 2^n and 3.
+        num_feat (int): Channel number of intermediate features.
+
+    """
+
+    def __init__(self, scale, num_feat, num_out_ch, input_resolution=None):
+        self.num_feat = num_feat
+        self.input_resolution = input_resolution
+        m = []
+        m.append(nn.Conv2d(num_feat, (scale**2) * num_out_ch, 3, 1, 1))
+        m.append(nn.PixelShuffle(scale))
+        super(UpsampleOneStep, self).__init__(*m)
+
+    def flops(self):
+        h, w = self.input_resolution
+        flops = h * w * self.num_feat * 3 * 9
+        return flops
+
diff --git a/ccrestoration/config/__init__.py b/ccrestoration/config/__init__.py
index 46a12bc..3f0900b 100644
--- a/ccrestoration/config/__init__.py
+++ b/ccrestoration/config/__init__.py
@@ -12,3 +12,4 @@
 from ccrestoration.config.iconvsr_config import IconVSRConfig  # noqa
 from ccrestoration.config.animesr_config import AnimeSRConfig  # noqa
 from ccrestoration.config.scunet_config import SCUNetConfig  # noqa
+from ccrestoration.config.dat_config import DATConfig  # noqa
\ No newline at end of file
diff --git a/ccrestoration/config/dat_config.py b/ccrestoration/config/dat_config.py
new file mode 100644
index 0000000..1b11228
--- /dev/null
+++ b/ccrestoration/config/dat_config.py
@@ -0,0 +1,172 @@
+from typing import Any, List, Optional, Tuple, Union
+
+from pydantic import field_validator
+from torch import nn
+
+from ccrestoration.config import CONFIG_REGISTRY
+from ccrestoration.type import ArchType, BaseConfig, ConfigType, ModelType
+
+class DATConfig(BaseConfig):
+    arch: Union[ArchType, str] = ArchType.DAT
+    model: Union[ModelType, str] = ModelType.DAT
+    scale: int = 2
+    in_chans: int = 3
+    img_size: Union[int, Tuple[int, ...]] = 64
+    img_range: float = 1.0
+    split_size: Union[List[int], Tuple[int, ...]] = [2, 4]  # noqa
+    depth: Union[List[int], Tuple[int, ...]] = [6, 6, 6, 6, 6, 6]  # noqa
+    embed_dim: int = 180
+    num_heads: Union[List[int], Tuple[int, ...]] = [6, 6, 6, 6, 6, 6] # noqa
+    expansion_factor: float = 4.0
+    resi_connection: str = "1conv"
+
+    qkv_bias: bool = True
+    qk_scale: Optional[float] = None
+    drop_rate: float = 0.0
+    attn_drop_rate: float = 0.0
+    drop_path_rate: float = 0.1
+    act_layer: Any = nn.GELU
+    norm_layer: Any = nn.LayerNorm
+    use_chk: bool = False
+    upsampler: str = "pixelshuffle"
+
+    @field_validator("scale")
+    def scale_match(cls, v: int) -> int:
+        if v not in [1, 2, 3, 4, 8]:
+            raise ValueError("scale factor must be one of [1, 2, 3, 4, 8]")
+        return v
+
+    @field_validator("upsampler")
+    def upsampler_match(cls, v: str) -> str:
+        if v not in ["pixelshuffle", "pixelshuffledirect", "nearest+conv", ""]:
+            raise ValueError("Upsampler must be one of ['pixelshuffle','pixelshuffledirect','nearest+conv', '']")
+        return v
+
+    @field_validator("resi_connection")
+    def resi_connection_match(cls, v: str) -> str:
+        if v not in ["1conv", "3conv"]:
+            raise ValueError("Residual connection must be one of ['1conv', '3conv']")
+        return v
+
+
+DATConfigs = [
+    # official models
+    # dat_s
+    DATConfig(
+        name=ConfigType.DAT_S_2x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_S_x2.pth",
+        hash="160330dd8a40b141e12713ca9bfde09e36a03c533b455965f157d023672cb794",
+        scale=2,
+        split_size=[8, 16],
+        expansion_factor=2,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_S_3x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_S_x3.pth",
+        hash="d1446d3eb2fbaad472c6fd6c2ac03a3467265c5f93822a741b09838b80f18b62",
+        scale=3,
+        split_size=[8, 16],
+        expansion_factor=2,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_S_4x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_S_x4.pth",
+        hash="2ba8a7cbe2fd88f3499d08d1b353fa17a98b815832f4426a7144a0ef9f3bfcf7",
+        scale=4,
+        split_size=[8, 16],
+        expansion_factor=2,
+    ),
+    # dat_m
+    DATConfig(
+        name=ConfigType.DAT_2x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_x2.pth",
+        hash="7760aa96e4ee77e29d4f89c3a4486200042e019461fdb8aa286f49aa00b89b51",
+        scale=2,
+        split_size=[8, 32],
+        expansion_factor=4,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_3x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_x3.pth",
+        hash="581973e02c06f90d4eb90acf743ec9604f56f3c2c6f9e1e2c2b38ded1f80d197",
+        scale=3,
+        split_size=[8, 32],
+        expansion_factor=4,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_4x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_x4.pth",
+        hash="391a6ce69899dff5ea3214557e9d585608254579217169faf3d4c353caff049e",
+        scale=4,
+        split_size=[8, 32],
+        expansion_factor=4,
+    ),
+    # dat_2
+    DATConfig(
+        name=ConfigType.DAT_2_2x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_2_x2.pth",
+        hash="aea2c76996c2b3e7cb034380e07738608bd59cc34c667331df7269e4b670ac18",
+        scale=2,
+        split_size=[8, 32],
+        expansion_factor=2,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_2_3x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_2_x3.pth",
+        hash="e19fbb2e6addf5cecf90472937fabdd905853b79c4ef807b4dd184c30bb22a28",
+        scale=3,
+        split_size=[8, 32],
+        expansion_factor=2,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_2_4x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_2_x4.pth",
+        hash="05b5c17bb5d1939ec0ec6b9368368d82d8c45b80c134e370f798efec0aeec395",
+        scale=4,
+        split_size=[8, 32],
+        expansion_factor=2,
+    ),
+    #dat_light
+    DATConfig(
+        name=ConfigType.DAT_light_2x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_light_x2.pth",
+        hash="3888bf067c1a3790adb64210e108690859f4900ffd158fe39a2cfe057ae8300c",
+        scale=2,
+        split_size=[8, 32],
+        expansion_factor=2,
+        depth=[18],
+        embed_dim=60,
+        num_heads=[6],
+        resi_connection='3conv',
+        upsampler='pixelshuffledirect'
+    ),
+    DATConfig(
+        name=ConfigType.DAT_light_3x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_light_x3.pth",
+        hash="dfb99deab865db05771080c73c50a06b67c9499c3a27ca64ab3729fa573d5cfe",
+        scale=3,
+        split_size=[8, 32],
+        expansion_factor=2,
+        depth=[18],
+        embed_dim=60,
+        num_heads=[6],
+        resi_connection='3conv',
+        upsampler='pixelshuffledirect'
+    ),
+    DATConfig(
+        name=ConfigType.DAT_light_4x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_light_x4.pth",
+        hash="e18223fc41500e72e6bd576f2a576eae0818f58fa8cc8646a2233f5faf7d0f74",
+        scale=4,
+        split_size=[8, 32],
+        expansion_factor=2,
+        depth=[18],
+        embed_dim=60,
+        num_heads=[6],
+        resi_connection='3conv',
+        upsampler='pixelshuffledirect'
+    )
+]
+
+for cfg in DATConfigs:
+    CONFIG_REGISTRY.register(cfg)
diff --git a/ccrestoration/model/__init__.py b/ccrestoration/model/__init__.py
index cd8a4ab..1b99c92 100644
--- a/ccrestoration/model/__init__.py
+++ b/ccrestoration/model/__init__.py
@@ -16,3 +16,4 @@
 from ccrestoration.model.iconvsr_model import IconVSRModel  # noqa
 from ccrestoration.model.animesr_model import AnimeSRModel  # noqa
 from ccrestoration.model.scunet_model import SCUNetModel  # noqa
+from ccrestoration.model.dat_model import DATModel  # noqa
diff --git a/ccrestoration/model/dat_model.py b/ccrestoration/model/dat_model.py
new file mode 100644
index 0000000..c3fb907
--- /dev/null
+++ b/ccrestoration/model/dat_model.py
@@ -0,0 +1,45 @@
+from typing import Any
+
+from ccrestoration.arch import DAT
+from ccrestoration.config import DATConfig
+from ccrestoration.model import MODEL_REGISTRY
+from ccrestoration.model.sr_base_model import SRBaseModel
+from ccrestoration.type import ModelType
+
+
+@MODEL_REGISTRY.register(name=ModelType.DAT)
+class DATModel(SRBaseModel):
+    def load_model(self) -> Any:
+        cfg: DATConfig = self.config
+        state_dict = self.get_state_dict()
+
+        if "params_ema" in state_dict:
+            state_dict = state_dict["params_ema"]
+        elif "params" in state_dict:
+            state_dict = state_dict["params"]
+
+        model = DAT(
+            img_size=cfg.img_size,
+            in_chans=cfg.in_chans,
+            embed_dim=cfg.embed_dim,
+            split_size=cfg.split_size,
+            depth=cfg.depth,
+            num_heads=cfg.num_heads,
+            qkv_bias=cfg.qkv_bias,
+            qk_scale=cfg.qk_scale,
+            drop_rate=cfg.drop_rate,
+            attn_drop_rate=cfg.attn_drop_rate,
+            drop_path_rate=cfg.drop_path_rate,
+            act_layer=cfg.act_layer,
+            norm_layer=cfg.norm_layer,
+            use_chk=cfg.use_chk,
+            upscale=cfg.scale,
+            img_range=cfg.img_range,
+            upsampler=cfg.upsampler,
+            resi_connection=cfg.resi_connection,
+            expansion_factor=cfg.expansion_factor
+        )
+
+        model.load_state_dict(state_dict)
+        model.eval().to(self.device)
+        return model
diff --git a/ccrestoration/type/arch.py b/ccrestoration/type/arch.py
index f49633f..dff1c6a 100644
--- a/ccrestoration/type/arch.py
+++ b/ccrestoration/type/arch.py
@@ -11,6 +11,7 @@ class ArchType(str, Enum):
     EDSR = "EDSR"
     SWINIR = "SWINIR"
     SCUNET = "SCUNET"
+    DAT = "DAT"
 
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 
diff --git a/ccrestoration/type/config.py b/ccrestoration/type/config.py
index 0d64345..c6ebff0 100644
--- a/ccrestoration/type/config.py
+++ b/ccrestoration/type/config.py
@@ -56,6 +56,20 @@ class ConfigType(str, Enum):
     SCUNet_color_real_psnr_1x = "SCUNet_color_real_psnr_1x.pth"
     SCUNet_color_real_gan_1x = "SCUNet_color_real_gan_1x.pth"
 
+    # DAT
+    DAT_S_2x = "DAT_S_x2.pth"
+    DAT_S_3x = "DAT_S_x3.pth"
+    DAT_S_4x = "DAT_S_x4.pth"
+    DAT_2x = "DAT_x2.pth"
+    DAT_3x = "DAT_x3.pth"
+    DAT_4x = "DAT_x4.pth"
+    DAT_2_2x = "DAT_2_x2.pth"
+    DAT_2_3x = "DAT_2_x3.pth"
+    DAT_2_4x = "DAT_2_x4.pth"
+    DAT_light_2x = "DAT_light_x2.pth"
+    DAT_light_3x = "DAT_light_x3.pth"
+    DAT_light_4x = "DAT_light_x4.pth"
+
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 
     # SpyNet
diff --git a/ccrestoration/type/model.py b/ccrestoration/type/model.py
index 9a842bf..0915d7f 100644
--- a/ccrestoration/type/model.py
+++ b/ccrestoration/type/model.py
@@ -10,6 +10,7 @@ class ModelType(str, Enum):
     EDSR = "EDSR"
     SwinIR = "SwinIR"
     SCUNet = "SCUNet"
+    DAT = "DAT"
 
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 
diff --git a/tests/test_dat.py b/tests/test_dat.py
new file mode 100644
index 0000000..17a5875
--- /dev/null
+++ b/tests/test_dat.py
@@ -0,0 +1,29 @@
+import cv2
+
+from ccrestoration import AutoConfig, AutoModel, BaseConfig, ConfigType
+from ccrestoration.model import SRBaseModel
+
+from .util import ASSETS_PATH, calculate_image_similarity, compare_image_size, get_device, load_image
+
+
+class Test_DAT:
+
+    def test_custom(self) -> None:
+        img1 = load_image()
+
+        for k in [
+            ConfigType.DAT_S_2x, ConfigType.DAT_S_3x, ConfigType.DAT_S_4x,
+            ConfigType.DAT_2x, ConfigType.DAT_3x, ConfigType.DAT_4x,
+            ConfigType.DAT_2_2x, ConfigType.DAT_2_3x, ConfigType.DAT_2_4x,
+            ConfigType.DAT_light_2x, ConfigType.DAT_light_3x, ConfigType.DAT_light_4x
+        ]:
+            print(f"Testing {k}")
+            cfg: BaseConfig = AutoConfig.from_pretrained(k)
+            model: SRBaseModel = AutoModel.from_config(config=cfg, fp16=False, device=get_device())
+            print(model.device)
+
+            img2 = model.inference_image(img1)
+            cv2.imwrite(str(ASSETS_PATH / f"test_{k}_out.jpg"), img2)
+
+            assert calculate_image_similarity(img1, img2)
+            assert compare_image_size(img1, img2, cfg.scale)

From 269f2ea097a02927e13a27310f24decf0c9c813d Mon Sep 17 00:00:00 2001
From: hyw-dev <1315176066@qq.com>
Date: Fri, 8 Nov 2024 17:59:45 +0800
Subject: [PATCH 2/5] feat: support DAT

---
 ccrestoration/arch/__init__.py     |    1 +
 ccrestoration/arch/dat_arch.py     | 1049 ++++++++++++++++++++++++++++
 ccrestoration/config/__init__.py   |    1 +
 ccrestoration/config/dat_config.py |  172 +++++
 ccrestoration/model/__init__.py    |    1 +
 ccrestoration/model/dat_model.py   |   45 ++
 ccrestoration/type/arch.py         |    1 +
 ccrestoration/type/config.py       |   14 +
 ccrestoration/type/model.py        |    1 +
 tests/test_dat.py                  |   29 +
 10 files changed, 1314 insertions(+)
 create mode 100644 ccrestoration/arch/dat_arch.py
 create mode 100644 ccrestoration/config/dat_config.py
 create mode 100644 ccrestoration/model/dat_model.py
 create mode 100644 tests/test_dat.py

diff --git a/ccrestoration/arch/__init__.py b/ccrestoration/arch/__init__.py
index 908ab3e..596d872 100644
--- a/ccrestoration/arch/__init__.py
+++ b/ccrestoration/arch/__init__.py
@@ -13,3 +13,4 @@
 from ccrestoration.arch.iconvsr_arch import IconVSR  # noqa
 from ccrestoration.arch.msrswvsr_arch import MSRSWVSR  # noqa
 from ccrestoration.arch.scunet_arch import SCUNet  # noqa
+from ccrestoration.arch.dat_arch import DAT  # noqa
diff --git a/ccrestoration/arch/dat_arch.py b/ccrestoration/arch/dat_arch.py
new file mode 100644
index 0000000..705a485
--- /dev/null
+++ b/ccrestoration/arch/dat_arch.py
@@ -0,0 +1,1049 @@
+# type: ignore
+import numpy as np
+import torch
+from einops import rearrange
+from einops.layers.torch import Rearrange
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.init import trunc_normal_
+from torch.utils import checkpoint
+
+from ccrestoration.arch import ARCH_REGISTRY
+from ccrestoration.arch.arch_util import Upsample, DropPath
+from ccrestoration.type import ArchType
+
+
+@ARCH_REGISTRY.register(name=ArchType.DAT)
+class DAT(nn.Module):
+    """Dual Aggregation Transformer
+    Args:
+        img_size (int): Input image size. Default: 64
+        in_chans (int): Number of input image channels. Default: 3
+        embed_dim (int): Patch embedding dimension. Default: 180
+        depths (tuple(int)): Depth of each residual group (number of DATB in each RG).
+        split_size (tuple(int)): Height and Width of spatial window.
+        num_heads (tuple(int)): Number of attention heads in different residual groups.
+        expansion_factor (float): Ratio of ffn hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        act_layer (nn.Module): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm
+        use_chk (bool): Whether to use checkpointing to save memory.
+        upscale: Upscale factor. 2/3/4 for image SR
+        img_range: Image range. 1. or 255.
+        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
+    """
+
+    def __init__(
+        self,
+        img_size=64,
+        in_chans=3,
+        embed_dim=180,
+        split_size=[2, 4],
+        depth=[2, 2, 2, 2],
+        num_heads=[2, 2, 2, 2],
+        expansion_factor=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.1,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        use_chk=False,
+        upscale=2,
+        img_range=1.0,
+        resi_connection="1conv",
+        upsampler="pixelshuffle",
+        **kwargs,
+    ):
+        super().__init__()
+
+        num_in_ch = in_chans
+        num_out_ch = in_chans
+        num_feat = 64
+        self.img_range = img_range
+        if in_chans == 3:
+            rgb_mean = (0.5, 0.5, 0.5)
+            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
+        else:
+            self.mean = torch.zeros(1, 1, 1, 1)
+        self.upscale = upscale
+        self.upsampler = upsampler
+
+        # ------------------------- 1, Shallow Feature Extraction ------------------------- #
+        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)
+
+        # ------------------------- 2, Deep Feature Extraction ------------------------- #
+        self.num_layers = len(depth)
+        self.use_chk = use_chk
+        self.num_features = self.embed_dim = (
+            embed_dim  # num_features for consistency with other models
+        )
+        heads = num_heads
+
+        self.before_RG = nn.Sequential(
+            Rearrange("b c h w -> b (h w) c"), nn.LayerNorm(embed_dim)
+        )
+
+        curr_dim = embed_dim
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, np.sum(depth))
+        ]  # stochastic depth decay rule
+
+        self.layers = nn.ModuleList()
+        for i in range(self.num_layers):
+            layer = ResidualGroup(
+                dim=embed_dim,
+                num_heads=heads[i],
+                reso=img_size,
+                split_size=split_size,
+                expansion_factor=expansion_factor,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_paths=dpr[sum(depth[:i]) : sum(depth[: i + 1])],
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                depth=depth[i],
+                use_chk=use_chk,
+                resi_connection=resi_connection,
+                rg_idx=i,
+            )
+            self.layers.append(layer)
+
+        self.norm = norm_layer(curr_dim)
+        # build the last conv layer in deep feature extraction
+        if resi_connection == "1conv":
+            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
+        elif resi_connection == "3conv":
+            # to save parameters and memory
+            self.conv_after_body = nn.Sequential(
+                nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1),
+            )
+
+        # ------------------------- 3, Reconstruction ------------------------- #
+        if self.upsampler == "pixelshuffle":
+            # for classical SR
+            self.conv_before_upsample = nn.Sequential(
+                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
+            )
+            self.upsample = Upsample(upscale, num_feat)
+            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+        elif self.upsampler == "pixelshuffledirect":
+            # for lightweight SR (to save parameters)
+            self.upsample = UpsampleOneStep(
+                upscale, embed_dim, num_out_ch, (img_size, img_size)
+            )
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(
+            m, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d)
+        ):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward_features(self, x):
+        _, _, H, W = x.shape
+        x_size = [H, W]
+        x = self.before_RG(x)
+        for layer in self.layers:
+            x = layer(x, x_size)
+        x = self.norm(x)
+        x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W)
+
+        return x
+
+    def forward(self, x):
+        """Input: x: (B, C, H, W)"""
+        self.mean = self.mean.type_as(x)
+        x = (x - self.mean) * self.img_range
+
+        if self.upsampler == "pixelshuffle":
+            # for image SR
+            x = self.conv_first(x)
+            x = self.conv_after_body(self.forward_features(x)) + x
+            x = self.conv_before_upsample(x)
+            x = self.conv_last(self.upsample(x))
+        elif self.upsampler == "pixelshuffledirect":
+            # for lightweight SR
+            x = self.conv_first(x)
+            x = self.conv_after_body(self.forward_features(x)) + x
+            x = self.upsample(x)
+
+        x = x / self.img_range + self.mean
+        return x
+
+def img2windows(img, H_sp, W_sp):
+    """Input: Image (B, C, H, W)
+    Output: Window Partition (B', N, C)
+    """
+    B, C, H, W = img.shape
+    img_reshape = img.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp)
+    img_perm = (
+        img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, H_sp * W_sp, C)
+    )
+    return img_perm
+
+
+def windows2img(img_splits_hw, H_sp, W_sp, H, W):
+    """Input: Window Partition (B', N, C)
+    Output: Image (B, H, W, C)
+    """
+    B = int(img_splits_hw.shape[0] / (H * W / H_sp / W_sp))
+
+    img = img_splits_hw.view(B, H // H_sp, W // W_sp, H_sp, W_sp, -1)
+    img = img.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return img
+
+
+class SpatialGate(nn.Module):
+    """Spatial-Gate.
+
+    Args:
+    ----
+        dim (int): Half of input channels.
+
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.conv = nn.Conv2d(
+            dim, dim, kernel_size=3, stride=1, padding=1, groups=dim
+        )  # DW Conv
+
+    def forward(self, x, H, W):
+        # Split
+        x1, x2 = x.chunk(2, dim=-1)
+        B, N, C = x.shape
+        x2 = (
+            self.conv(self.norm(x2).transpose(1, 2).contiguous().view(B, C // 2, H, W))
+            .flatten(2)
+            .transpose(-1, -2)
+            .contiguous()
+        )
+
+        return x1 * x2
+
+
+class SGFN(nn.Module):
+    """Spatial-Gate Feed-Forward Network.
+
+    Args:
+    ----
+        in_features (int): Number of input channels.
+        hidden_features (int | None): Number of hidden channels. Default: None
+        out_features (int | None): Number of output channels. Default: None
+        act_layer (nn.Module): Activation layer. Default: nn.GELU
+        drop (float): Dropout rate. Default: 0.0
+
+    """
+
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.sg = SpatialGate(hidden_features // 2)
+        self.fc2 = nn.Linear(hidden_features // 2, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x, H, W):
+        """Input: x: (B, H*W, C), H, W
+        Output: x: (B, H*W, C)
+        """
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+
+        x = self.sg(x, H, W)
+        x = self.drop(x)
+
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class DynamicPosBias(nn.Module):
+    # The implementation builds on Crossformer code https://github.com/cheerss/CrossFormer/blob/main/models/crossformer.py
+
+    """Dynamic Relative Position Bias.
+
+    Args:
+    ----
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        residual (bool):  If True, use residual strage to connect conv.
+
+    """
+
+    def __init__(self, dim, num_heads, residual):
+        super().__init__()
+        self.residual = residual
+        self.num_heads = num_heads
+        self.pos_dim = dim // 4
+        self.pos_proj = nn.Linear(2, self.pos_dim)
+        self.pos1 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.pos_dim),
+        )
+        self.pos2 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.pos_dim),
+        )
+        self.pos3 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.num_heads),
+        )
+
+    def forward(self, biases):
+        if self.residual:
+            pos = self.pos_proj(biases)  # 2Gh-1 * 2Gw-1, heads
+            pos = pos + self.pos1(pos)
+            pos = pos + self.pos2(pos)
+            pos = self.pos3(pos)
+        else:
+            pos = self.pos3(self.pos2(self.pos1(self.pos_proj(biases))))
+        return pos
+
+
+class Spatial_Attention(nn.Module):
+    """Spatial Window Self-Attention.
+    It supports rectangle window (containing square window).
+
+    Args:
+    ----
+        dim (int): Number of input channels.
+        idx (int): The indentix of window. (0/1)
+        split_size (tuple(int)): Height and Width of spatial window.
+        dim_out (int | None): The dimension of the attention output. Default: None
+        num_heads (int): Number of attention heads. Default: 6
+        attn_drop (float): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float): Dropout ratio of output. Default: 0.0
+        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set
+        position_bias (bool): The dynamic relative position bias. Default: True
+
+    """
+
+    def __init__(
+        self,
+        dim,
+        idx,
+        split_size=[8, 8],
+        dim_out=None,
+        num_heads=6,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        qk_scale=None,
+        position_bias=True,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.dim_out = dim_out or dim
+        self.split_size = split_size
+        self.num_heads = num_heads
+        self.idx = idx
+        self.position_bias = position_bias
+
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        if idx == 0:
+            H_sp, W_sp = self.split_size[0], self.split_size[1]
+        elif idx == 1:
+            W_sp, H_sp = self.split_size[0], self.split_size[1]
+        else:
+            print("ERROR MODE", idx)
+            exit(0)
+        self.H_sp = H_sp
+        self.W_sp = W_sp
+
+        if self.position_bias:
+            self.pos = DynamicPosBias(self.dim // 4, self.num_heads, residual=False)
+            # generate mother-set
+            position_bias_h = torch.arange(1 - self.H_sp, self.H_sp)
+            position_bias_w = torch.arange(1 - self.W_sp, self.W_sp)
+            biases = torch.stack(
+                torch.meshgrid([position_bias_h, position_bias_w], indexing="ij")
+            )
+            biases = biases.flatten(1).transpose(0, 1).contiguous().float()
+            self.register_buffer("rpe_biases", biases)
+
+            # get pair-wise relative position index for each token inside the window
+            coords_h = torch.arange(self.H_sp)
+            coords_w = torch.arange(self.W_sp)
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))
+            coords_flatten = torch.flatten(coords, 1)
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+            relative_coords[:, :, 0] += self.H_sp - 1
+            relative_coords[:, :, 1] += self.W_sp - 1
+            relative_coords[:, :, 0] *= 2 * self.W_sp - 1
+            relative_position_index = relative_coords.sum(-1)
+            self.register_buffer("relative_position_index", relative_position_index)
+
+        self.attn_drop = nn.Dropout(attn_drop)
+
+    def im2win(self, x, H, W):
+        B, N, C = x.shape
+        x = x.transpose(-2, -1).contiguous().view(B, C, H, W)
+        x = img2windows(x, self.H_sp, self.W_sp)
+        x = (
+            x.reshape(-1, self.H_sp * self.W_sp, self.num_heads, C // self.num_heads)
+            .permute(0, 2, 1, 3)
+            .contiguous()
+        )
+        return x
+
+    def forward(self, qkv, H, W, mask=None):
+        """Input: qkv: (B, 3*L, C), H, W, mask: (B, N, N), N is the window size
+        Output: x (B, H, W, C)
+        """
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        B, L, C = q.shape
+        assert L == H * W, "flatten img_tokens has wrong size"
+
+        # partition the q,k,v, image to window
+        q = self.im2win(q, H, W)
+        k = self.im2win(k, H, W)
+        v = self.im2win(v, H, W)
+
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)  # B head N C @ B head C N --> B head N N
+
+        # calculate drpe
+        if self.position_bias:
+            pos = self.pos(self.rpe_biases)
+            # select position bias
+            relative_position_bias = pos[self.relative_position_index.view(-1)].view(
+                self.H_sp * self.W_sp, self.H_sp * self.W_sp, -1
+            )
+            relative_position_bias = relative_position_bias.permute(
+                2, 0, 1
+            ).contiguous()
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+        N = attn.shape[3]
+
+        # use mask for shift window
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(
+                0
+            )
+            attn = attn.view(-1, self.num_heads, N, N)
+
+        attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype)
+        attn = self.attn_drop(attn)
+
+        x = attn @ v
+        x = x.transpose(1, 2).reshape(
+            -1, self.H_sp * self.W_sp, C
+        )  # B head N N @ B head N C
+
+        # merge the window, window to image
+        x = windows2img(x, self.H_sp, self.W_sp, H, W)  # B H' W' C
+
+        return x
+
+
+class Axial_Spatial_Attention(nn.Module):
+    """Axial Spatial Self-Attention
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads. Default: 6
+        split_size (tuple(int)): Height and Width of spatial window.
+        shift_size (tuple(int)): Shift size for spatial window.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float): Dropout rate. Default: 0.0
+        attn_drop (float): Attention dropout rate. Default: 0.0
+        rg_idx (int): The indentix of Residual Group (RG)
+        b_idx (int): The indentix of Block in each RG
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        reso=64,
+        split_size=[8, 8],
+        shift_size=[1, 2],
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        rg_idx=0,
+        b_idx=0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.split_size = split_size
+        self.shift_size = shift_size
+        self.b_idx = b_idx
+        self.rg_idx = rg_idx
+        self.patches_resolution = reso
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+
+        assert (
+            0 <= self.shift_size[0] < self.split_size[0]
+        ), "shift_size must in 0-split_size0"
+        assert (
+            0 <= self.shift_size[1] < self.split_size[1]
+        ), "shift_size must in 0-split_size1"
+
+        self.branch_num = 2
+
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(drop)
+
+        self.attns = nn.ModuleList([
+            Spatial_Attention(
+                dim // 2,
+                idx=i,
+                split_size=split_size,
+                num_heads=num_heads // 2,
+                dim_out=dim // 2,
+                qk_scale=qk_scale,
+                attn_drop=attn_drop,
+                proj_drop=drop,
+                position_bias=True,
+            )
+            for i in range(self.branch_num)
+        ])
+
+        if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or (
+            self.rg_idx % 2 != 0 and self.b_idx % 4 == 0
+        ):
+            attn_mask = self.calculate_mask(
+                self.patches_resolution, self.patches_resolution
+            )
+            self.register_buffer("attn_mask_0", attn_mask[0])
+            self.register_buffer("attn_mask_1", attn_mask[1])
+        else:
+            attn_mask = None
+            self.register_buffer("attn_mask_0", None)
+            self.register_buffer("attn_mask_1", None)
+
+        self.dwconv = nn.Sequential(
+            nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim),
+            nn.BatchNorm2d(dim),
+            nn.GELU(),
+        )
+        self.channel_interaction = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(dim, dim // 8, kernel_size=1),
+            nn.BatchNorm2d(dim // 8),
+            nn.GELU(),
+            nn.Conv2d(dim // 8, dim, kernel_size=1),
+        )
+        self.spatial_interaction = nn.Sequential(
+            nn.Conv2d(dim, dim // 16, kernel_size=1),
+            nn.BatchNorm2d(dim // 16),
+            nn.GELU(),
+            nn.Conv2d(dim // 16, 1, kernel_size=1),
+        )
+
+    def calculate_mask(self, H, W):
+        # The implementation builds on Swin Transformer code https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
+        # calculate attention mask for shift window
+        img_mask_0 = torch.zeros((1, H, W, 1))  # 1 H W 1 idx=0
+        img_mask_1 = torch.zeros((1, H, W, 1))  # 1 H W 1 idx=1
+        h_slices_0 = (
+            slice(0, -self.split_size[0]),
+            slice(-self.split_size[0], -self.shift_size[0]),
+            slice(-self.shift_size[0], None),
+        )
+        w_slices_0 = (
+            slice(0, -self.split_size[1]),
+            slice(-self.split_size[1], -self.shift_size[1]),
+            slice(-self.shift_size[1], None),
+        )
+
+        h_slices_1 = (
+            slice(0, -self.split_size[1]),
+            slice(-self.split_size[1], -self.shift_size[1]),
+            slice(-self.shift_size[1], None),
+        )
+        w_slices_1 = (
+            slice(0, -self.split_size[0]),
+            slice(-self.split_size[0], -self.shift_size[0]),
+            slice(-self.shift_size[0], None),
+        )
+        cnt = 0
+        for h in h_slices_0:
+            for w in w_slices_0:
+                img_mask_0[:, h, w, :] = cnt
+                cnt += 1
+        cnt = 0
+        for h in h_slices_1:
+            for w in w_slices_1:
+                img_mask_1[:, h, w, :] = cnt
+                cnt += 1
+
+        # calculate mask for window-0
+        img_mask_0 = img_mask_0.view(
+            1,
+            H // self.split_size[0],
+            self.split_size[0],
+            W // self.split_size[1],
+            self.split_size[1],
+            1,
+        )
+        img_mask_0 = (
+            img_mask_0.permute(0, 1, 3, 2, 4, 5)
+            .contiguous()
+            .view(-1, self.split_size[0], self.split_size[1], 1)
+        )  # nW, sw[0], sw[1], 1
+        mask_windows_0 = img_mask_0.view(-1, self.split_size[0] * self.split_size[1])
+        attn_mask_0 = mask_windows_0.unsqueeze(1) - mask_windows_0.unsqueeze(2)
+        attn_mask_0 = attn_mask_0.masked_fill(attn_mask_0 != 0, -100.0).masked_fill(
+            attn_mask_0 == 0, 0.0
+        )
+
+        # calculate mask for window-1
+        img_mask_1 = img_mask_1.view(
+            1,
+            H // self.split_size[1],
+            self.split_size[1],
+            W // self.split_size[0],
+            self.split_size[0],
+            1,
+        )
+        img_mask_1 = (
+            img_mask_1.permute(0, 1, 3, 2, 4, 5)
+            .contiguous()
+            .view(-1, self.split_size[1], self.split_size[0], 1)
+        )  # nW, sw[1], sw[0], 1
+        mask_windows_1 = img_mask_1.view(-1, self.split_size[1] * self.split_size[0])
+        attn_mask_1 = mask_windows_1.unsqueeze(1) - mask_windows_1.unsqueeze(2)
+        attn_mask_1 = attn_mask_1.masked_fill(attn_mask_1 != 0, -100.0).masked_fill(
+            attn_mask_1 == 0, 0.0
+        )
+
+        return attn_mask_0, attn_mask_1
+
+    def forward(self, x, H, W):
+        """Input: x: (B, H*W, C), H, W
+        Output: x: (B, H*W, C)
+        """
+        B, L, C = x.shape
+        assert L == H * W, "flatten img_tokens has wrong size"
+
+        qkv = self.qkv(x).reshape(B, -1, 3, C).permute(2, 0, 1, 3)  # 3, B, HW, C
+        # V without partition
+        v = qkv[2].transpose(-2, -1).contiguous().view(B, C, H, W)
+
+        # image padding
+        max_split_size = max(self.split_size[0], self.split_size[1])
+        pad_l = pad_t = 0
+        pad_r = (max_split_size - W % max_split_size) % max_split_size
+        pad_b = (max_split_size - H % max_split_size) % max_split_size
+
+        qkv = qkv.reshape(3 * B, H, W, C).permute(0, 3, 1, 2)  # 3B C H W
+        qkv = (
+            F.pad(qkv, (pad_l, pad_r, pad_t, pad_b))
+            .reshape(3, B, C, -1)
+            .transpose(-2, -1)
+        )  # l r t b
+        _H = pad_b + H
+        _W = pad_r + W
+        _L = _H * _W
+
+        # window-0 and window-1 on split channels [C/2, C/2]; for square windows (e.g., 8x8), window-0 and window-1 can be merged
+        # shift in block: (0, 4, 8, ...), (2, 6, 10, ...), (0, 4, 8, ...), (2, 6, 10, ...), ...
+        if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or (
+            self.rg_idx % 2 != 0 and self.b_idx % 4 == 0
+        ):
+            qkv = qkv.view(3, B, _H, _W, C)
+            qkv_0 = torch.roll(
+                qkv[:, :, :, :, : C // 2],
+                shifts=(-self.shift_size[0], -self.shift_size[1]),
+                dims=(2, 3),
+            )
+            qkv_0 = qkv_0.view(3, B, _L, C // 2)
+            qkv_1 = torch.roll(
+                qkv[:, :, :, :, C // 2 :],
+                shifts=(-self.shift_size[1], -self.shift_size[0]),
+                dims=(2, 3),
+            )
+            qkv_1 = qkv_1.view(3, B, _L, C // 2)
+
+            if self.patches_resolution != _H or self.patches_resolution != _W:
+                mask_tmp = self.calculate_mask(_H, _W)
+                x1_shift = self.attns[0](qkv_0, _H, _W, mask=mask_tmp[0].to(x.device))
+                x2_shift = self.attns[1](qkv_1, _H, _W, mask=mask_tmp[1].to(x.device))
+            else:
+                x1_shift = self.attns[0](qkv_0, _H, _W, mask=self.attn_mask_0)
+                x2_shift = self.attns[1](qkv_1, _H, _W, mask=self.attn_mask_1)
+
+            x1 = torch.roll(
+                x1_shift, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2)
+            )
+            x2 = torch.roll(
+                x2_shift, shifts=(self.shift_size[1], self.shift_size[0]), dims=(1, 2)
+            )
+            x1 = x1[:, :H, :W, :].reshape(B, L, C // 2)
+            x2 = x2[:, :H, :W, :].reshape(B, L, C // 2)
+            # attention output
+            attened_x = torch.cat([x1, x2], dim=2)
+
+        else:
+            x1 = self.attns[0](qkv[:, :, :, : C // 2], _H, _W)[:, :H, :W, :].reshape(
+                B, L, C // 2
+            )
+            x2 = self.attns[1](qkv[:, :, :, C // 2 :], _H, _W)[:, :H, :W, :].reshape(
+                B, L, C // 2
+            )
+            # attention output
+            attened_x = torch.cat([x1, x2], dim=2)
+
+        # convolution output
+        conv_x = self.dwconv(v)
+
+        # Adaptive Interaction Module (AIM)
+        # C-Map (before sigmoid)
+        channel_map = (
+            self.channel_interaction(conv_x)
+            .permute(0, 2, 3, 1)
+            .contiguous()
+            .view(B, 1, C)
+        )
+        # S-Map (before sigmoid)
+        attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W)
+        spatial_map = self.spatial_interaction(attention_reshape)
+
+        # C-I
+        attened_x = attened_x * torch.sigmoid(channel_map)
+        # S-I
+        conv_x = torch.sigmoid(spatial_map) * conv_x
+        conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(B, L, C)
+
+        x = attened_x + conv_x
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class Axial_Channel_Attention(nn.Module):
+    # The implementation builds on XCiT code https://github.com/facebookresearch/xcit
+
+    """Axial Channel Self-Attention
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads. Default: 6
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set.
+        attn_drop (float): Attention dropout rate. Default: 0.0
+        drop_path (float): Stochastic depth rate. Default: 0.0
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.dwconv = nn.Sequential(
+            nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim),
+            nn.BatchNorm2d(dim),
+            nn.GELU(),
+        )
+        self.channel_interaction = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(dim, dim // 8, kernel_size=1),
+            nn.BatchNorm2d(dim // 8),
+            nn.GELU(),
+            nn.Conv2d(dim // 8, dim, kernel_size=1),
+        )
+        self.spatial_interaction = nn.Sequential(
+            nn.Conv2d(dim, dim // 16, kernel_size=1),
+            nn.BatchNorm2d(dim // 16),
+            nn.GELU(),
+            nn.Conv2d(dim // 16, 1, kernel_size=1),
+        )
+
+    def forward(self, x, H, W):
+        """Input: x: (B, H*W, C), H, W
+        Output: x: (B, H*W, C)
+        """
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        q = q.transpose(-2, -1)
+        k = k.transpose(-2, -1)
+        v = v.transpose(-2, -1)
+
+        v_ = v.reshape(B, C, N).contiguous().view(B, C, H, W)
+
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        # attention output
+        attened_x = (attn @ v).permute(0, 3, 1, 2).reshape(B, N, C)
+
+        # convolution output
+        conv_x = self.dwconv(v_)
+
+        # Adaptive Interaction Module (AIM)
+        # C-Map (before sigmoid)
+        attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W)
+        channel_map = self.channel_interaction(attention_reshape)
+        # S-Map (before sigmoid)
+        spatial_map = (
+            self.spatial_interaction(conv_x)
+            .permute(0, 2, 3, 1)
+            .contiguous()
+            .view(B, N, 1)
+        )
+
+        # S-I
+        attened_x = attened_x * torch.sigmoid(spatial_map)
+        # C-I
+        conv_x = conv_x * torch.sigmoid(channel_map)
+        conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(B, N, C)
+
+        x = attened_x + conv_x
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class DATB(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        reso=64,
+        split_size=[2, 4],
+        shift_size=[1, 2],
+        expansion_factor=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        rg_idx=0,
+        b_idx=0,
+    ):
+        super().__init__()
+
+        self.norm1 = norm_layer(dim)
+
+        if b_idx % 2 == 0:
+            # DSTB
+            self.attn = Axial_Spatial_Attention(
+                dim,
+                num_heads=num_heads,
+                reso=reso,
+                split_size=split_size,
+                shift_size=shift_size,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop,
+                attn_drop=attn_drop,
+                rg_idx=rg_idx,
+                b_idx=b_idx,
+            )
+        else:
+            # DCTB
+            self.attn = Axial_Channel_Attention(
+                dim,
+                num_heads=num_heads,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                attn_drop=attn_drop,
+                proj_drop=drop,
+            )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        ffn_hidden_dim = int(dim * expansion_factor)
+        self.ffn = SGFN(
+            in_features=dim,
+            hidden_features=ffn_hidden_dim,
+            out_features=dim,
+            act_layer=act_layer,
+        )
+        self.norm2 = norm_layer(dim)
+
+    def forward(self, x, x_size):
+        """Input: x: (B, H*W, C), x_size: (H, W)
+        Output: x: (B, H*W, C)
+        """
+        H, W = x_size
+        x = x + self.drop_path(self.attn(self.norm1(x), H, W))
+        x = x + self.drop_path(self.ffn(self.norm2(x), H, W))
+
+        return x
+
+
+class ResidualGroup(nn.Module):
+    """ResidualGroup
+    Args:
+        dim (int): Number of input channels.
+        reso (int): Input resolution.
+        num_heads (int): Number of attention heads.
+        split_size (tuple(int)): Height and Width of spatial window.
+        expansion_factor (float): Ratio of ffn hidden dim to embedding dim.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop (float): Dropout rate. Default: 0
+        attn_drop(float): Attention dropout rate. Default: 0
+        drop_paths (float | None): Stochastic depth rate.
+        act_layer (nn.Module): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm
+        depth (int): Number of dual aggregation Transformer blocks in residual group.
+        use_chk (bool): Whether to use checkpointing to save memory.
+        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
+    """
+
+    def __init__(
+        self,
+        dim,
+        reso,
+        num_heads,
+        split_size=[2, 4],
+        expansion_factor=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_paths=None,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        depth=2,
+        use_chk=False,
+        resi_connection="1conv",
+        rg_idx=0,
+    ):
+        super().__init__()
+        self.use_chk = use_chk
+        self.reso = reso
+
+        self.blocks = nn.ModuleList([
+            DATB(
+                dim=dim,
+                num_heads=num_heads,
+                reso=reso,
+                split_size=split_size,
+                shift_size=[split_size[0] // 2, split_size[1] // 2],
+                expansion_factor=expansion_factor,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop,
+                attn_drop=attn_drop,
+                drop_path=drop_paths[i],
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                rg_idx=rg_idx,
+                b_idx=i,
+            )
+            for i in range(depth)
+        ])
+
+        if resi_connection == "1conv":
+            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
+        elif resi_connection == "3conv":
+            self.conv = nn.Sequential(
+                nn.Conv2d(dim, dim // 4, 3, 1, 1),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(dim // 4, dim, 3, 1, 1),
+            )
+
+    def forward(self, x, x_size):
+        """Input: x: (B, H*W, C), x_size: (H, W)
+        Output: x: (B, H*W, C)
+        """
+        H, W = x_size
+        res = x
+        for blk in self.blocks:
+            if self.use_chk:
+                x = checkpoint.checkpoint(blk, x, x_size, use_reentrant=False)
+            else:
+                x = blk(x, x_size)
+        x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W)
+        x = self.conv(x)
+        x = rearrange(x, "b c h w -> b (h w) c")
+        x = res + x
+
+        return x
+
+
+class UpsampleOneStep(nn.Sequential):
+    """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
+       Used in lightweight SR to save parameters.
+
+    Args:
+    ----
+        scale (int): Scale factor. Supported scales: 2^n and 3.
+        num_feat (int): Channel number of intermediate features.
+
+    """
+
+    def __init__(self, scale, num_feat, num_out_ch, input_resolution=None):
+        self.num_feat = num_feat
+        self.input_resolution = input_resolution
+        m = []
+        m.append(nn.Conv2d(num_feat, (scale**2) * num_out_ch, 3, 1, 1))
+        m.append(nn.PixelShuffle(scale))
+        super(UpsampleOneStep, self).__init__(*m)
+
+    def flops(self):
+        h, w = self.input_resolution
+        flops = h * w * self.num_feat * 3 * 9
+        return flops
+
diff --git a/ccrestoration/config/__init__.py b/ccrestoration/config/__init__.py
index 46a12bc..3f0900b 100644
--- a/ccrestoration/config/__init__.py
+++ b/ccrestoration/config/__init__.py
@@ -12,3 +12,4 @@
 from ccrestoration.config.iconvsr_config import IconVSRConfig  # noqa
 from ccrestoration.config.animesr_config import AnimeSRConfig  # noqa
 from ccrestoration.config.scunet_config import SCUNetConfig  # noqa
+from ccrestoration.config.dat_config import DATConfig  # noqa
\ No newline at end of file
diff --git a/ccrestoration/config/dat_config.py b/ccrestoration/config/dat_config.py
new file mode 100644
index 0000000..1b11228
--- /dev/null
+++ b/ccrestoration/config/dat_config.py
@@ -0,0 +1,172 @@
+from typing import Any, List, Optional, Tuple, Union
+
+from pydantic import field_validator
+from torch import nn
+
+from ccrestoration.config import CONFIG_REGISTRY
+from ccrestoration.type import ArchType, BaseConfig, ConfigType, ModelType
+
+class DATConfig(BaseConfig):
+    arch: Union[ArchType, str] = ArchType.DAT
+    model: Union[ModelType, str] = ModelType.DAT
+    scale: int = 2
+    in_chans: int = 3
+    img_size: Union[int, Tuple[int, ...]] = 64
+    img_range: float = 1.0
+    split_size: Union[List[int], Tuple[int, ...]] = [2, 4]  # noqa
+    depth: Union[List[int], Tuple[int, ...]] = [6, 6, 6, 6, 6, 6]  # noqa
+    embed_dim: int = 180
+    num_heads: Union[List[int], Tuple[int, ...]] = [6, 6, 6, 6, 6, 6] # noqa
+    expansion_factor: float = 4.0
+    resi_connection: str = "1conv"
+
+    qkv_bias: bool = True
+    qk_scale: Optional[float] = None
+    drop_rate: float = 0.0
+    attn_drop_rate: float = 0.0
+    drop_path_rate: float = 0.1
+    act_layer: Any = nn.GELU
+    norm_layer: Any = nn.LayerNorm
+    use_chk: bool = False
+    upsampler: str = "pixelshuffle"
+
+    @field_validator("scale")
+    def scale_match(cls, v: int) -> int:
+        if v not in [1, 2, 3, 4, 8]:
+            raise ValueError("scale factor must be one of [1, 2, 3, 4, 8]")
+        return v
+
+    @field_validator("upsampler")
+    def upsampler_match(cls, v: str) -> str:
+        if v not in ["pixelshuffle", "pixelshuffledirect", "nearest+conv", ""]:
+            raise ValueError("Upsampler must be one of ['pixelshuffle','pixelshuffledirect','nearest+conv', '']")
+        return v
+
+    @field_validator("resi_connection")
+    def resi_connection_match(cls, v: str) -> str:
+        if v not in ["1conv", "3conv"]:
+            raise ValueError("Residual connection must be one of ['1conv', '3conv']")
+        return v
+
+
+DATConfigs = [
+    # official models
+    # dat_s
+    DATConfig(
+        name=ConfigType.DAT_S_2x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_S_x2.pth",
+        hash="160330dd8a40b141e12713ca9bfde09e36a03c533b455965f157d023672cb794",
+        scale=2,
+        split_size=[8, 16],
+        expansion_factor=2,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_S_3x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_S_x3.pth",
+        hash="d1446d3eb2fbaad472c6fd6c2ac03a3467265c5f93822a741b09838b80f18b62",
+        scale=3,
+        split_size=[8, 16],
+        expansion_factor=2,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_S_4x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_S_x4.pth",
+        hash="2ba8a7cbe2fd88f3499d08d1b353fa17a98b815832f4426a7144a0ef9f3bfcf7",
+        scale=4,
+        split_size=[8, 16],
+        expansion_factor=2,
+    ),
+    # dat_m
+    DATConfig(
+        name=ConfigType.DAT_2x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_x2.pth",
+        hash="7760aa96e4ee77e29d4f89c3a4486200042e019461fdb8aa286f49aa00b89b51",
+        scale=2,
+        split_size=[8, 32],
+        expansion_factor=4,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_3x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_x3.pth",
+        hash="581973e02c06f90d4eb90acf743ec9604f56f3c2c6f9e1e2c2b38ded1f80d197",
+        scale=3,
+        split_size=[8, 32],
+        expansion_factor=4,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_4x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_x4.pth",
+        hash="391a6ce69899dff5ea3214557e9d585608254579217169faf3d4c353caff049e",
+        scale=4,
+        split_size=[8, 32],
+        expansion_factor=4,
+    ),
+    # dat_2
+    DATConfig(
+        name=ConfigType.DAT_2_2x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_2_x2.pth",
+        hash="aea2c76996c2b3e7cb034380e07738608bd59cc34c667331df7269e4b670ac18",
+        scale=2,
+        split_size=[8, 32],
+        expansion_factor=2,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_2_3x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_2_x3.pth",
+        hash="e19fbb2e6addf5cecf90472937fabdd905853b79c4ef807b4dd184c30bb22a28",
+        scale=3,
+        split_size=[8, 32],
+        expansion_factor=2,
+    ),
+    DATConfig(
+        name=ConfigType.DAT_2_4x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_2_x4.pth",
+        hash="05b5c17bb5d1939ec0ec6b9368368d82d8c45b80c134e370f798efec0aeec395",
+        scale=4,
+        split_size=[8, 32],
+        expansion_factor=2,
+    ),
+    #dat_light
+    DATConfig(
+        name=ConfigType.DAT_light_2x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_light_x2.pth",
+        hash="3888bf067c1a3790adb64210e108690859f4900ffd158fe39a2cfe057ae8300c",
+        scale=2,
+        split_size=[8, 32],
+        expansion_factor=2,
+        depth=[18],
+        embed_dim=60,
+        num_heads=[6],
+        resi_connection='3conv',
+        upsampler='pixelshuffledirect'
+    ),
+    DATConfig(
+        name=ConfigType.DAT_light_3x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_light_x3.pth",
+        hash="dfb99deab865db05771080c73c50a06b67c9499c3a27ca64ab3729fa573d5cfe",
+        scale=3,
+        split_size=[8, 32],
+        expansion_factor=2,
+        depth=[18],
+        embed_dim=60,
+        num_heads=[6],
+        resi_connection='3conv',
+        upsampler='pixelshuffledirect'
+    ),
+    DATConfig(
+        name=ConfigType.DAT_light_4x,
+        url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_light_x4.pth",
+        hash="e18223fc41500e72e6bd576f2a576eae0818f58fa8cc8646a2233f5faf7d0f74",
+        scale=4,
+        split_size=[8, 32],
+        expansion_factor=2,
+        depth=[18],
+        embed_dim=60,
+        num_heads=[6],
+        resi_connection='3conv',
+        upsampler='pixelshuffledirect'
+    )
+]
+
+for cfg in DATConfigs:
+    CONFIG_REGISTRY.register(cfg)
diff --git a/ccrestoration/model/__init__.py b/ccrestoration/model/__init__.py
index cd8a4ab..1b99c92 100644
--- a/ccrestoration/model/__init__.py
+++ b/ccrestoration/model/__init__.py
@@ -16,3 +16,4 @@
 from ccrestoration.model.iconvsr_model import IconVSRModel  # noqa
 from ccrestoration.model.animesr_model import AnimeSRModel  # noqa
 from ccrestoration.model.scunet_model import SCUNetModel  # noqa
+from ccrestoration.model.dat_model import DATModel  # noqa
diff --git a/ccrestoration/model/dat_model.py b/ccrestoration/model/dat_model.py
new file mode 100644
index 0000000..c3fb907
--- /dev/null
+++ b/ccrestoration/model/dat_model.py
@@ -0,0 +1,45 @@
+from typing import Any
+
+from ccrestoration.arch import DAT
+from ccrestoration.config import DATConfig
+from ccrestoration.model import MODEL_REGISTRY
+from ccrestoration.model.sr_base_model import SRBaseModel
+from ccrestoration.type import ModelType
+
+
+@MODEL_REGISTRY.register(name=ModelType.DAT)
+class DATModel(SRBaseModel):
+    def load_model(self) -> Any:
+        cfg: DATConfig = self.config
+        state_dict = self.get_state_dict()
+
+        if "params_ema" in state_dict:
+            state_dict = state_dict["params_ema"]
+        elif "params" in state_dict:
+            state_dict = state_dict["params"]
+
+        model = DAT(
+            img_size=cfg.img_size,
+            in_chans=cfg.in_chans,
+            embed_dim=cfg.embed_dim,
+            split_size=cfg.split_size,
+            depth=cfg.depth,
+            num_heads=cfg.num_heads,
+            qkv_bias=cfg.qkv_bias,
+            qk_scale=cfg.qk_scale,
+            drop_rate=cfg.drop_rate,
+            attn_drop_rate=cfg.attn_drop_rate,
+            drop_path_rate=cfg.drop_path_rate,
+            act_layer=cfg.act_layer,
+            norm_layer=cfg.norm_layer,
+            use_chk=cfg.use_chk,
+            upscale=cfg.scale,
+            img_range=cfg.img_range,
+            upsampler=cfg.upsampler,
+            resi_connection=cfg.resi_connection,
+            expansion_factor=cfg.expansion_factor
+        )
+
+        model.load_state_dict(state_dict)
+        model.eval().to(self.device)
+        return model
diff --git a/ccrestoration/type/arch.py b/ccrestoration/type/arch.py
index f49633f..dff1c6a 100644
--- a/ccrestoration/type/arch.py
+++ b/ccrestoration/type/arch.py
@@ -11,6 +11,7 @@ class ArchType(str, Enum):
     EDSR = "EDSR"
     SWINIR = "SWINIR"
     SCUNET = "SCUNET"
+    DAT = "DAT"
 
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 
diff --git a/ccrestoration/type/config.py b/ccrestoration/type/config.py
index 0d64345..c6ebff0 100644
--- a/ccrestoration/type/config.py
+++ b/ccrestoration/type/config.py
@@ -56,6 +56,20 @@ class ConfigType(str, Enum):
     SCUNet_color_real_psnr_1x = "SCUNet_color_real_psnr_1x.pth"
     SCUNet_color_real_gan_1x = "SCUNet_color_real_gan_1x.pth"
 
+    # DAT
+    DAT_S_2x = "DAT_S_x2.pth"
+    DAT_S_3x = "DAT_S_x3.pth"
+    DAT_S_4x = "DAT_S_x4.pth"
+    DAT_2x = "DAT_x2.pth"
+    DAT_3x = "DAT_x3.pth"
+    DAT_4x = "DAT_x4.pth"
+    DAT_2_2x = "DAT_2_x2.pth"
+    DAT_2_3x = "DAT_2_x3.pth"
+    DAT_2_4x = "DAT_2_x4.pth"
+    DAT_light_2x = "DAT_light_x2.pth"
+    DAT_light_3x = "DAT_light_x3.pth"
+    DAT_light_4x = "DAT_light_x4.pth"
+
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 
     # SpyNet
diff --git a/ccrestoration/type/model.py b/ccrestoration/type/model.py
index 9a842bf..0915d7f 100644
--- a/ccrestoration/type/model.py
+++ b/ccrestoration/type/model.py
@@ -10,6 +10,7 @@ class ModelType(str, Enum):
     EDSR = "EDSR"
     SwinIR = "SwinIR"
     SCUNet = "SCUNet"
+    DAT = "DAT"
 
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 
diff --git a/tests/test_dat.py b/tests/test_dat.py
new file mode 100644
index 0000000..17a5875
--- /dev/null
+++ b/tests/test_dat.py
@@ -0,0 +1,29 @@
+import cv2
+
+from ccrestoration import AutoConfig, AutoModel, BaseConfig, ConfigType
+from ccrestoration.model import SRBaseModel
+
+from .util import ASSETS_PATH, calculate_image_similarity, compare_image_size, get_device, load_image
+
+
+class Test_DAT:
+
+    def test_custom(self) -> None:
+        img1 = load_image()
+
+        for k in [
+            ConfigType.DAT_S_2x, ConfigType.DAT_S_3x, ConfigType.DAT_S_4x,
+            ConfigType.DAT_2x, ConfigType.DAT_3x, ConfigType.DAT_4x,
+            ConfigType.DAT_2_2x, ConfigType.DAT_2_3x, ConfigType.DAT_2_4x,
+            ConfigType.DAT_light_2x, ConfigType.DAT_light_3x, ConfigType.DAT_light_4x
+        ]:
+            print(f"Testing {k}")
+            cfg: BaseConfig = AutoConfig.from_pretrained(k)
+            model: SRBaseModel = AutoModel.from_config(config=cfg, fp16=False, device=get_device())
+            print(model.device)
+
+            img2 = model.inference_image(img1)
+            cv2.imwrite(str(ASSETS_PATH / f"test_{k}_out.jpg"), img2)
+
+            assert calculate_image_similarity(img1, img2)
+            assert compare_image_size(img1, img2, cfg.scale)

From d71df2032ee61dd59dc339b779905f1861b64ad4 Mon Sep 17 00:00:00 2001
From: hyw-dev <1315176066@qq.com>
Date: Fri, 8 Nov 2024 18:21:49 +0800
Subject: [PATCH 3/5] fix: Do not use mutable data structures for argument
 defaults

---
 ccrestoration/arch/dat_arch.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/ccrestoration/arch/dat_arch.py b/ccrestoration/arch/dat_arch.py
index 705a485..2aeae58 100644
--- a/ccrestoration/arch/dat_arch.py
+++ b/ccrestoration/arch/dat_arch.py
@@ -42,9 +42,9 @@ def __init__(
         img_size=64,
         in_chans=3,
         embed_dim=180,
-        split_size=[2, 4],
-        depth=[2, 2, 2, 2],
-        num_heads=[2, 2, 2, 2],
+        split_size=[2, 4], # noqa
+        depth=[2, 2, 2, 2], # noqa
+        num_heads=[2, 2, 2, 2], # noqa
         expansion_factor=4.0,
         qkv_bias=True,
         qk_scale=None,
@@ -355,7 +355,7 @@ def __init__(
         self,
         dim,
         idx,
-        split_size=[8, 8],
+        split_size=[8, 8],  # noqa
         dim_out=None,
         num_heads=6,
         attn_drop=0.0,
@@ -494,8 +494,8 @@ def __init__(
         dim,
         num_heads,
         reso=64,
-        split_size=[8, 8],
-        shift_size=[1, 2],
+        split_size=[8, 8], # noqa
+        shift_size=[1, 2], # noqa
         qkv_bias=False,
         qk_scale=None,
         drop=0.0,
@@ -865,8 +865,8 @@ def __init__(
         dim,
         num_heads,
         reso=64,
-        split_size=[2, 4],
-        shift_size=[1, 2],
+        split_size=[2, 4], # noqa
+        shift_size=[1, 2], # noqa
         expansion_factor=4.0,
         qkv_bias=False,
         qk_scale=None,
@@ -954,7 +954,7 @@ def __init__(
         dim,
         reso,
         num_heads,
-        split_size=[2, 4],
+        split_size=[2, 4], # noqa
         expansion_factor=4.0,
         qkv_bias=False,
         qk_scale=None,

From fdf6a3a65f6419c0466deb32c828d1322a281810 Mon Sep 17 00:00:00 2001
From: hyw-dev <1315176066@qq.com>
Date: Fri, 8 Nov 2024 18:21:49 +0800
Subject: [PATCH 4/5] fix: Do not use mutable data structures for argument
 defaults

---
 ccrestoration/arch/dat_arch.py     | 208 +++++++++++------------------
 ccrestoration/config/__init__.py   |   2 +-
 ccrestoration/config/dat_config.py |  19 +--
 ccrestoration/model/dat_model.py   |   2 +-
 tests/test_dat.py                  |  17 ++-
 5 files changed, 99 insertions(+), 149 deletions(-)

diff --git a/ccrestoration/arch/dat_arch.py b/ccrestoration/arch/dat_arch.py
index 705a485..a199975 100644
--- a/ccrestoration/arch/dat_arch.py
+++ b/ccrestoration/arch/dat_arch.py
@@ -9,7 +9,7 @@
 from torch.utils import checkpoint
 
 from ccrestoration.arch import ARCH_REGISTRY
-from ccrestoration.arch.arch_util import Upsample, DropPath
+from ccrestoration.arch.arch_util import DropPath, Upsample
 from ccrestoration.type import ArchType
 
 
@@ -42,9 +42,9 @@ def __init__(
         img_size=64,
         in_chans=3,
         embed_dim=180,
-        split_size=[2, 4],
-        depth=[2, 2, 2, 2],
-        num_heads=[2, 2, 2, 2],
+        split_size=[2, 4],  # noqa
+        depth=[2, 2, 2, 2],  # noqa
+        num_heads=[2, 2, 2, 2],  # noqa
         expansion_factor=4.0,
         qkv_bias=True,
         qk_scale=None,
@@ -80,19 +80,13 @@ def __init__(
         # ------------------------- 2, Deep Feature Extraction ------------------------- #
         self.num_layers = len(depth)
         self.use_chk = use_chk
-        self.num_features = self.embed_dim = (
-            embed_dim  # num_features for consistency with other models
-        )
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
         heads = num_heads
 
-        self.before_RG = nn.Sequential(
-            Rearrange("b c h w -> b (h w) c"), nn.LayerNorm(embed_dim)
-        )
+        self.before_RG = nn.Sequential(Rearrange("b c h w -> b (h w) c"), nn.LayerNorm(embed_dim))
 
         curr_dim = embed_dim
-        dpr = [
-            x.item() for x in torch.linspace(0, drop_path_rate, np.sum(depth))
-        ]  # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, np.sum(depth))]  # stochastic depth decay rule
 
         self.layers = nn.ModuleList()
         for i in range(self.num_layers):
@@ -140,9 +134,7 @@ def __init__(
             self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
         elif self.upsampler == "pixelshuffledirect":
             # for lightweight SR (to save parameters)
-            self.upsample = UpsampleOneStep(
-                upscale, embed_dim, num_out_ch, (img_size, img_size)
-            )
+            self.upsample = UpsampleOneStep(upscale, embed_dim, num_out_ch, (img_size, img_size))
 
         self.apply(self._init_weights)
 
@@ -151,9 +143,7 @@ def _init_weights(self, m):
             trunc_normal_(m.weight, std=0.02)
             if isinstance(m, nn.Linear) and m.bias is not None:
                 nn.init.constant_(m.bias, 0)
-        elif isinstance(
-            m, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d)
-        ):
+        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d)):
             nn.init.constant_(m.bias, 0)
             nn.init.constant_(m.weight, 1.0)
 
@@ -188,15 +178,14 @@ def forward(self, x):
         x = x / self.img_range + self.mean
         return x
 
+
 def img2windows(img, H_sp, W_sp):
     """Input: Image (B, C, H, W)
     Output: Window Partition (B', N, C)
     """
     B, C, H, W = img.shape
     img_reshape = img.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp)
-    img_perm = (
-        img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, H_sp * W_sp, C)
-    )
+    img_perm = img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, H_sp * W_sp, C)
     return img_perm
 
 
@@ -223,9 +212,7 @@ class SpatialGate(nn.Module):
     def __init__(self, dim):
         super().__init__()
         self.norm = nn.LayerNorm(dim)
-        self.conv = nn.Conv2d(
-            dim, dim, kernel_size=3, stride=1, padding=1, groups=dim
-        )  # DW Conv
+        self.conv = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim)  # DW Conv
 
     def forward(self, x, H, W):
         # Split
@@ -355,7 +342,7 @@ def __init__(
         self,
         dim,
         idx,
-        split_size=[8, 8],
+        split_size=[8, 8],  # noqa
         dim_out=None,
         num_heads=6,
         attn_drop=0.0,
@@ -389,9 +376,7 @@ def __init__(
             # generate mother-set
             position_bias_h = torch.arange(1 - self.H_sp, self.H_sp)
             position_bias_w = torch.arange(1 - self.W_sp, self.W_sp)
-            biases = torch.stack(
-                torch.meshgrid([position_bias_h, position_bias_w], indexing="ij")
-            )
+            biases = torch.stack(torch.meshgrid([position_bias_h, position_bias_w], indexing="ij"))
             biases = biases.flatten(1).transpose(0, 1).contiguous().float()
             self.register_buffer("rpe_biases", biases)
 
@@ -414,11 +399,7 @@ def im2win(self, x, H, W):
         B, N, C = x.shape
         x = x.transpose(-2, -1).contiguous().view(B, C, H, W)
         x = img2windows(x, self.H_sp, self.W_sp)
-        x = (
-            x.reshape(-1, self.H_sp * self.W_sp, self.num_heads, C // self.num_heads)
-            .permute(0, 2, 1, 3)
-            .contiguous()
-        )
+        x = x.reshape(-1, self.H_sp * self.W_sp, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).contiguous()
         return x
 
     def forward(self, qkv, H, W, mask=None):
@@ -445,9 +426,7 @@ def forward(self, qkv, H, W, mask=None):
             relative_position_bias = pos[self.relative_position_index.view(-1)].view(
                 self.H_sp * self.W_sp, self.H_sp * self.W_sp, -1
             )
-            relative_position_bias = relative_position_bias.permute(
-                2, 0, 1
-            ).contiguous()
+            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
             attn = attn + relative_position_bias.unsqueeze(0)
 
         N = attn.shape[3]
@@ -455,18 +434,14 @@ def forward(self, qkv, H, W, mask=None):
         # use mask for shift window
         if mask is not None:
             nW = mask.shape[0]
-            attn = attn.view(B, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(
-                0
-            )
+            attn = attn.view(B, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
             attn = attn.view(-1, self.num_heads, N, N)
 
         attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype)
         attn = self.attn_drop(attn)
 
         x = attn @ v
-        x = x.transpose(1, 2).reshape(
-            -1, self.H_sp * self.W_sp, C
-        )  # B head N N @ B head N C
+        x = x.transpose(1, 2).reshape(-1, self.H_sp * self.W_sp, C)  # B head N N @ B head N C
 
         # merge the window, window to image
         x = windows2img(x, self.H_sp, self.W_sp, H, W)  # B H' W' C
@@ -494,8 +469,8 @@ def __init__(
         dim,
         num_heads,
         reso=64,
-        split_size=[8, 8],
-        shift_size=[1, 2],
+        split_size=[8, 8],  # noqa
+        shift_size=[1, 2],  # noqa
         qkv_bias=False,
         qk_scale=None,
         drop=0.0,
@@ -513,39 +488,35 @@ def __init__(
         self.patches_resolution = reso
         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
 
-        assert (
-            0 <= self.shift_size[0] < self.split_size[0]
-        ), "shift_size must in 0-split_size0"
-        assert (
-            0 <= self.shift_size[1] < self.split_size[1]
-        ), "shift_size must in 0-split_size1"
+        assert 0 <= self.shift_size[0] < self.split_size[0], "shift_size must in 0-split_size0"
+        assert 0 <= self.shift_size[1] < self.split_size[1], "shift_size must in 0-split_size1"
 
         self.branch_num = 2
 
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(drop)
 
-        self.attns = nn.ModuleList([
-            Spatial_Attention(
-                dim // 2,
-                idx=i,
-                split_size=split_size,
-                num_heads=num_heads // 2,
-                dim_out=dim // 2,
-                qk_scale=qk_scale,
-                attn_drop=attn_drop,
-                proj_drop=drop,
-                position_bias=True,
-            )
-            for i in range(self.branch_num)
-        ])
+        self.attns = nn.ModuleList(
+            [
+                Spatial_Attention(
+                    dim // 2,
+                    idx=i,
+                    split_size=split_size,
+                    num_heads=num_heads // 2,
+                    dim_out=dim // 2,
+                    qk_scale=qk_scale,
+                    attn_drop=attn_drop,
+                    proj_drop=drop,
+                    position_bias=True,
+                )
+                for i in range(self.branch_num)
+            ]
+        )
 
         if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or (
             self.rg_idx % 2 != 0 and self.b_idx % 4 == 0
         ):
-            attn_mask = self.calculate_mask(
-                self.patches_resolution, self.patches_resolution
-            )
+            attn_mask = self.calculate_mask(self.patches_resolution, self.patches_resolution)
             self.register_buffer("attn_mask_0", attn_mask[0])
             self.register_buffer("attn_mask_1", attn_mask[1])
         else:
@@ -619,15 +590,11 @@ def calculate_mask(self, H, W):
             1,
         )
         img_mask_0 = (
-            img_mask_0.permute(0, 1, 3, 2, 4, 5)
-            .contiguous()
-            .view(-1, self.split_size[0], self.split_size[1], 1)
+            img_mask_0.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.split_size[0], self.split_size[1], 1)
         )  # nW, sw[0], sw[1], 1
         mask_windows_0 = img_mask_0.view(-1, self.split_size[0] * self.split_size[1])
         attn_mask_0 = mask_windows_0.unsqueeze(1) - mask_windows_0.unsqueeze(2)
-        attn_mask_0 = attn_mask_0.masked_fill(attn_mask_0 != 0, -100.0).masked_fill(
-            attn_mask_0 == 0, 0.0
-        )
+        attn_mask_0 = attn_mask_0.masked_fill(attn_mask_0 != 0, -100.0).masked_fill(attn_mask_0 == 0, 0.0)
 
         # calculate mask for window-1
         img_mask_1 = img_mask_1.view(
@@ -639,15 +606,11 @@ def calculate_mask(self, H, W):
             1,
         )
         img_mask_1 = (
-            img_mask_1.permute(0, 1, 3, 2, 4, 5)
-            .contiguous()
-            .view(-1, self.split_size[1], self.split_size[0], 1)
+            img_mask_1.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.split_size[1], self.split_size[0], 1)
         )  # nW, sw[1], sw[0], 1
         mask_windows_1 = img_mask_1.view(-1, self.split_size[1] * self.split_size[0])
         attn_mask_1 = mask_windows_1.unsqueeze(1) - mask_windows_1.unsqueeze(2)
-        attn_mask_1 = attn_mask_1.masked_fill(attn_mask_1 != 0, -100.0).masked_fill(
-            attn_mask_1 == 0, 0.0
-        )
+        attn_mask_1 = attn_mask_1.masked_fill(attn_mask_1 != 0, -100.0).masked_fill(attn_mask_1 == 0, 0.0)
 
         return attn_mask_0, attn_mask_1
 
@@ -669,11 +632,7 @@ def forward(self, x, H, W):
         pad_b = (max_split_size - H % max_split_size) % max_split_size
 
         qkv = qkv.reshape(3 * B, H, W, C).permute(0, 3, 1, 2)  # 3B C H W
-        qkv = (
-            F.pad(qkv, (pad_l, pad_r, pad_t, pad_b))
-            .reshape(3, B, C, -1)
-            .transpose(-2, -1)
-        )  # l r t b
+        qkv = F.pad(qkv, (pad_l, pad_r, pad_t, pad_b)).reshape(3, B, C, -1).transpose(-2, -1)  # l r t b
         _H = pad_b + H
         _W = pad_r + W
         _L = _H * _W
@@ -705,24 +664,16 @@ def forward(self, x, H, W):
                 x1_shift = self.attns[0](qkv_0, _H, _W, mask=self.attn_mask_0)
                 x2_shift = self.attns[1](qkv_1, _H, _W, mask=self.attn_mask_1)
 
-            x1 = torch.roll(
-                x1_shift, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2)
-            )
-            x2 = torch.roll(
-                x2_shift, shifts=(self.shift_size[1], self.shift_size[0]), dims=(1, 2)
-            )
+            x1 = torch.roll(x1_shift, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2))
+            x2 = torch.roll(x2_shift, shifts=(self.shift_size[1], self.shift_size[0]), dims=(1, 2))
             x1 = x1[:, :H, :W, :].reshape(B, L, C // 2)
             x2 = x2[:, :H, :W, :].reshape(B, L, C // 2)
             # attention output
             attened_x = torch.cat([x1, x2], dim=2)
 
         else:
-            x1 = self.attns[0](qkv[:, :, :, : C // 2], _H, _W)[:, :H, :W, :].reshape(
-                B, L, C // 2
-            )
-            x2 = self.attns[1](qkv[:, :, :, C // 2 :], _H, _W)[:, :H, :W, :].reshape(
-                B, L, C // 2
-            )
+            x1 = self.attns[0](qkv[:, :, :, : C // 2], _H, _W)[:, :H, :W, :].reshape(B, L, C // 2)
+            x2 = self.attns[1](qkv[:, :, :, C // 2 :], _H, _W)[:, :H, :W, :].reshape(B, L, C // 2)
             # attention output
             attened_x = torch.cat([x1, x2], dim=2)
 
@@ -731,12 +682,7 @@ def forward(self, x, H, W):
 
         # Adaptive Interaction Module (AIM)
         # C-Map (before sigmoid)
-        channel_map = (
-            self.channel_interaction(conv_x)
-            .permute(0, 2, 3, 1)
-            .contiguous()
-            .view(B, 1, C)
-        )
+        channel_map = self.channel_interaction(conv_x).permute(0, 2, 3, 1).contiguous().view(B, 1, C)
         # S-Map (before sigmoid)
         attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W)
         spatial_map = self.spatial_interaction(attention_reshape)
@@ -838,12 +784,7 @@ def forward(self, x, H, W):
         attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W)
         channel_map = self.channel_interaction(attention_reshape)
         # S-Map (before sigmoid)
-        spatial_map = (
-            self.spatial_interaction(conv_x)
-            .permute(0, 2, 3, 1)
-            .contiguous()
-            .view(B, N, 1)
-        )
+        spatial_map = self.spatial_interaction(conv_x).permute(0, 2, 3, 1).contiguous().view(B, N, 1)
 
         # S-I
         attened_x = attened_x * torch.sigmoid(spatial_map)
@@ -865,8 +806,8 @@ def __init__(
         dim,
         num_heads,
         reso=64,
-        split_size=[2, 4],
-        shift_size=[1, 2],
+        split_size=[2, 4],  # noqa
+        shift_size=[1, 2],  # noqa
         expansion_factor=4.0,
         qkv_bias=False,
         qk_scale=None,
@@ -954,7 +895,7 @@ def __init__(
         dim,
         reso,
         num_heads,
-        split_size=[2, 4],
+        split_size=[2, 4],  # noqa
         expansion_factor=4.0,
         qkv_bias=False,
         qk_scale=None,
@@ -972,26 +913,28 @@ def __init__(
         self.use_chk = use_chk
         self.reso = reso
 
-        self.blocks = nn.ModuleList([
-            DATB(
-                dim=dim,
-                num_heads=num_heads,
-                reso=reso,
-                split_size=split_size,
-                shift_size=[split_size[0] // 2, split_size[1] // 2],
-                expansion_factor=expansion_factor,
-                qkv_bias=qkv_bias,
-                qk_scale=qk_scale,
-                drop=drop,
-                attn_drop=attn_drop,
-                drop_path=drop_paths[i],
-                act_layer=act_layer,
-                norm_layer=norm_layer,
-                rg_idx=rg_idx,
-                b_idx=i,
-            )
-            for i in range(depth)
-        ])
+        self.blocks = nn.ModuleList(
+            [
+                DATB(
+                    dim=dim,
+                    num_heads=num_heads,
+                    reso=reso,
+                    split_size=split_size,
+                    shift_size=[split_size[0] // 2, split_size[1] // 2],
+                    expansion_factor=expansion_factor,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop,
+                    attn_drop=attn_drop,
+                    drop_path=drop_paths[i],
+                    act_layer=act_layer,
+                    norm_layer=norm_layer,
+                    rg_idx=rg_idx,
+                    b_idx=i,
+                )
+                for i in range(depth)
+            ]
+        )
 
         if resi_connection == "1conv":
             self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
@@ -1046,4 +989,3 @@ def flops(self):
         h, w = self.input_resolution
         flops = h * w * self.num_feat * 3 * 9
         return flops
-
diff --git a/ccrestoration/config/__init__.py b/ccrestoration/config/__init__.py
index 3f0900b..c67ead7 100644
--- a/ccrestoration/config/__init__.py
+++ b/ccrestoration/config/__init__.py
@@ -12,4 +12,4 @@
 from ccrestoration.config.iconvsr_config import IconVSRConfig  # noqa
 from ccrestoration.config.animesr_config import AnimeSRConfig  # noqa
 from ccrestoration.config.scunet_config import SCUNetConfig  # noqa
-from ccrestoration.config.dat_config import DATConfig  # noqa
\ No newline at end of file
+from ccrestoration.config.dat_config import DATConfig  # noqa
diff --git a/ccrestoration/config/dat_config.py b/ccrestoration/config/dat_config.py
index 1b11228..d49bde4 100644
--- a/ccrestoration/config/dat_config.py
+++ b/ccrestoration/config/dat_config.py
@@ -6,6 +6,7 @@
 from ccrestoration.config import CONFIG_REGISTRY
 from ccrestoration.type import ArchType, BaseConfig, ConfigType, ModelType
 
+
 class DATConfig(BaseConfig):
     arch: Union[ArchType, str] = ArchType.DAT
     model: Union[ModelType, str] = ModelType.DAT
@@ -16,7 +17,7 @@ class DATConfig(BaseConfig):
     split_size: Union[List[int], Tuple[int, ...]] = [2, 4]  # noqa
     depth: Union[List[int], Tuple[int, ...]] = [6, 6, 6, 6, 6, 6]  # noqa
     embed_dim: int = 180
-    num_heads: Union[List[int], Tuple[int, ...]] = [6, 6, 6, 6, 6, 6] # noqa
+    num_heads: Union[List[int], Tuple[int, ...]] = [6, 6, 6, 6, 6, 6]  # noqa
     expansion_factor: float = 4.0
     resi_connection: str = "1conv"
 
@@ -126,7 +127,7 @@ def resi_connection_match(cls, v: str) -> str:
         split_size=[8, 32],
         expansion_factor=2,
     ),
-    #dat_light
+    # dat_light
     DATConfig(
         name=ConfigType.DAT_light_2x,
         url="https://github.com/routineLife1/save_weights/releases/download/weights/DAT_light_x2.pth",
@@ -137,8 +138,8 @@ def resi_connection_match(cls, v: str) -> str:
         depth=[18],
         embed_dim=60,
         num_heads=[6],
-        resi_connection='3conv',
-        upsampler='pixelshuffledirect'
+        resi_connection="3conv",
+        upsampler="pixelshuffledirect",
     ),
     DATConfig(
         name=ConfigType.DAT_light_3x,
@@ -150,8 +151,8 @@ def resi_connection_match(cls, v: str) -> str:
         depth=[18],
         embed_dim=60,
         num_heads=[6],
-        resi_connection='3conv',
-        upsampler='pixelshuffledirect'
+        resi_connection="3conv",
+        upsampler="pixelshuffledirect",
     ),
     DATConfig(
         name=ConfigType.DAT_light_4x,
@@ -163,9 +164,9 @@ def resi_connection_match(cls, v: str) -> str:
         depth=[18],
         embed_dim=60,
         num_heads=[6],
-        resi_connection='3conv',
-        upsampler='pixelshuffledirect'
-    )
+        resi_connection="3conv",
+        upsampler="pixelshuffledirect",
+    ),
 ]
 
 for cfg in DATConfigs:
diff --git a/ccrestoration/model/dat_model.py b/ccrestoration/model/dat_model.py
index c3fb907..ca5315c 100644
--- a/ccrestoration/model/dat_model.py
+++ b/ccrestoration/model/dat_model.py
@@ -37,7 +37,7 @@ def load_model(self) -> Any:
             img_range=cfg.img_range,
             upsampler=cfg.upsampler,
             resi_connection=cfg.resi_connection,
-            expansion_factor=cfg.expansion_factor
+            expansion_factor=cfg.expansion_factor,
         )
 
         model.load_state_dict(state_dict)
diff --git a/tests/test_dat.py b/tests/test_dat.py
index 17a5875..53823a4 100644
--- a/tests/test_dat.py
+++ b/tests/test_dat.py
@@ -7,15 +7,22 @@
 
 
 class Test_DAT:
-
     def test_custom(self) -> None:
         img1 = load_image()
 
         for k in [
-            ConfigType.DAT_S_2x, ConfigType.DAT_S_3x, ConfigType.DAT_S_4x,
-            ConfigType.DAT_2x, ConfigType.DAT_3x, ConfigType.DAT_4x,
-            ConfigType.DAT_2_2x, ConfigType.DAT_2_3x, ConfigType.DAT_2_4x,
-            ConfigType.DAT_light_2x, ConfigType.DAT_light_3x, ConfigType.DAT_light_4x
+            ConfigType.DAT_S_2x,
+            ConfigType.DAT_S_3x,
+            ConfigType.DAT_S_4x,
+            ConfigType.DAT_2x,
+            ConfigType.DAT_3x,
+            ConfigType.DAT_4x,
+            ConfigType.DAT_2_2x,
+            ConfigType.DAT_2_3x,
+            ConfigType.DAT_2_4x,
+            ConfigType.DAT_light_2x,
+            ConfigType.DAT_light_3x,
+            ConfigType.DAT_light_4x,
         ]:
             print(f"Testing {k}")
             cfg: BaseConfig = AutoConfig.from_pretrained(k)

From a151fc8574b28a4d0d8ac888312f89e1e85a541a Mon Sep 17 00:00:00 2001
From: hyw-dev <1315176066@qq.com>
Date: Thu, 12 Dec 2024 13:19:59 +0800
Subject: [PATCH 5/5] feat: support HAT

---
 ccrestoration/arch/__init__.py     |    1 +
 ccrestoration/arch/hat_arch.py     | 1063 ++++++++++++++++++++++++++++
 ccrestoration/config/__init__.py   |    1 +
 ccrestoration/config/hat_config.py |  334 +++++++++
 ccrestoration/model/__init__.py    |    1 +
 ccrestoration/model/hat_model.py   |   53 ++
 ccrestoration/type/arch.py         |    1 +
 ccrestoration/type/config.py       |   16 +
 ccrestoration/type/model.py        |    1 +
 tests/test_hat.py                  |   76 ++
 10 files changed, 1547 insertions(+)
 create mode 100644 ccrestoration/arch/hat_arch.py
 create mode 100644 ccrestoration/config/hat_config.py
 create mode 100644 ccrestoration/model/hat_model.py
 create mode 100644 tests/test_hat.py

diff --git a/ccrestoration/arch/__init__.py b/ccrestoration/arch/__init__.py
index b55e0a6..be57604 100644
--- a/ccrestoration/arch/__init__.py
+++ b/ccrestoration/arch/__init__.py
@@ -15,3 +15,4 @@
 from ccrestoration.arch.scunet_arch import SCUNet  # noqa
 from ccrestoration.arch.dat_arch import DAT  # noqa
 from ccrestoration.arch.srcnn_arch import SRCNN  # noqa
+from ccrestoration.arch.hat_arch import HAT  # noqa
diff --git a/ccrestoration/arch/hat_arch.py b/ccrestoration/arch/hat_arch.py
new file mode 100644
index 0000000..19355dd
--- /dev/null
+++ b/ccrestoration/arch/hat_arch.py
@@ -0,0 +1,1063 @@
+# type: ignore
+import math
+
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn.init import trunc_normal_
+
+from ccrestoration.arch import ARCH_REGISTRY
+from ccrestoration.arch.arch_util import DropPath, to_2tuple
+from ccrestoration.type import ArchType
+
+
+@ARCH_REGISTRY.register(name=ArchType.HAT)
+class HAT(nn.Module):
+    r"""Hybrid Attention Transformer
+        A PyTorch implementation of : `Activating More Pixels in Image Super-Resolution Transformer`.
+        Some codes are based on SwinIR.
+
+    Args:
+    ----
+        img_size (int | tuple(int)): Input image size. Default 64
+        patch_size (int | tuple(int)): Patch size. Default: 1
+        in_chans (int): Number of input image channels. Default: 3
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depth (tuple(int)): Depth of each Swin Transformer layer.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        window_size (int): Window size. Default: 7
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True
+        upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
+        img_range: Image range. 1. or 255.
+        upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
+        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
+
+    """
+
+    def __init__(
+        self,
+        img_size=64,
+        patch_size=1,
+        in_chans=3,
+        embed_dim=96,
+        depth=(6, 6, 6, 6),
+        num_heads=(6, 6, 6, 6),
+        window_size=7,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.1,
+        norm_layer=nn.LayerNorm,
+        ape=False,
+        patch_norm=True,
+        upscale=2,
+        img_range=1.0,
+        upsampler="",
+        resi_connection="1conv",
+        **kwargs,
+    ):
+        super(HAT, self).__init__()
+
+        self.window_size = window_size
+        self.shift_size = window_size // 2
+        self.overlap_ratio = overlap_ratio
+
+        num_in_ch = in_chans
+        num_out_ch = in_chans
+        num_feat = 64
+        self.img_range = img_range
+        if in_chans == 3:
+            rgb_mean = (0.5, 0.5, 0.5)
+            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
+        else:
+            self.mean = torch.zeros(1, 1, 1, 1)
+        self.upscale = upscale
+        self.upsampler = upsampler
+
+        # relative position index
+        relative_position_index_SA = self.calculate_rpi_sa()
+        relative_position_index_OCA = self.calculate_rpi_oca()
+        self.register_buffer("relative_position_index_SA", relative_position_index_SA)
+        self.register_buffer("relative_position_index_OCA", relative_position_index_OCA)
+
+        # ------------------------- 1, shallow feature extraction ------------------------- #
+        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)
+
+        # ------------------------- 2, deep feature extraction ------------------------- #
+        self.num_layers = len(depth)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = embed_dim
+        self.mlp_ratio = mlp_ratio
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=embed_dim,
+            embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None,
+        )
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        # merge non-overlapping patches into image
+        self.patch_unembed = PatchUnEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=embed_dim,
+            embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None,
+        )
+
+        # absolute position embedding
+        if self.ape is True:
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+            trunc_normal_(self.absolute_pos_embed, std=0.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depth))]  # stochastic depth decay rule
+
+        # build Residual Hybrid Attention Groups (RHAG)
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = RHAG(
+                dim=embed_dim,
+                input_resolution=(patches_resolution[0], patches_resolution[1]),
+                depth=depth[i_layer],
+                num_heads=num_heads[i_layer],
+                window_size=window_size,
+                compress_ratio=compress_ratio,
+                squeeze_factor=squeeze_factor,
+                conv_scale=conv_scale,
+                overlap_ratio=overlap_ratio,
+                mlp_ratio=self.mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[sum(depth[:i_layer]) : sum(depth[: i_layer + 1])],  # no impact on SR results
+                norm_layer=norm_layer,
+                downsample=None,
+                img_size=img_size,
+                patch_size=patch_size,
+                resi_connection=resi_connection,
+            )
+            self.layers.append(layer)
+        self.norm = norm_layer(self.num_features)
+
+        # build the last conv layer in deep feature extraction
+        if resi_connection == "1conv":
+            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
+        elif resi_connection == "identity":
+            self.conv_after_body = nn.Identity()
+
+        # ------------------------- 3, high quality image reconstruction ------------------------- #
+        if self.upsampler == "pixelshuffle":
+            # for classical SR
+            self.conv_before_upsample = nn.Sequential(
+                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
+            )
+            self.upsample = Upsample(upscale, num_feat)
+            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def calculate_rpi_sa(self):
+        # calculate relative position index for SA
+        coords_h = torch.arange(self.window_size)
+        coords_w = torch.arange(self.window_size)
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        return relative_position_index
+
+    def calculate_rpi_oca(self):
+        # calculate relative position index for OCA
+        window_size_ori = self.window_size
+        window_size_ext = self.window_size + int(self.overlap_ratio * self.window_size)
+
+        coords_h = torch.arange(window_size_ori)
+        coords_w = torch.arange(window_size_ori)
+        coords_ori = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))  # 2, ws, ws
+        coords_ori_flatten = torch.flatten(coords_ori, 1)  # 2, ws*ws
+
+        coords_h = torch.arange(window_size_ext)
+        coords_w = torch.arange(window_size_ext)
+        coords_ext = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))  # 2, wse, wse
+        coords_ext_flatten = torch.flatten(coords_ext, 1)  # 2, wse*wse
+
+        relative_coords = coords_ext_flatten[:, None, :] - coords_ori_flatten[:, :, None]  # 2, ws*ws, wse*wse
+
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # ws*ws, wse*wse, 2
+        relative_coords[:, :, 0] += window_size_ori - window_size_ext + 1  # shift to start from 0
+        relative_coords[:, :, 1] += window_size_ori - window_size_ext + 1
+
+        relative_coords[:, :, 0] *= window_size_ori + window_size_ext - 1
+        relative_position_index = relative_coords.sum(-1)
+        return relative_position_index
+
+    def calculate_mask(self, x_size):
+        # calculate attention mask for SW-MSA
+        h, w = x_size
+        img_mask = torch.zeros((1, h, w, 1))  # 1 h w 1
+        h_slices = (
+            slice(0, -self.window_size),
+            slice(-self.window_size, -self.shift_size),
+            slice(-self.shift_size, None),
+        )
+        w_slices = (
+            slice(0, -self.window_size),
+            slice(-self.window_size, -self.shift_size),
+            slice(-self.shift_size, None),
+        )
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        mask_windows = window_partition(img_mask, self.window_size)  # nw, window_size, window_size, 1
+        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0).masked_fill(attn_mask == 0, 0.0)
+
+        return attn_mask
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"absolute_pos_embed"}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {"relative_position_bias_table"}
+
+    def forward_features(self, x):
+        x_size = (x.shape[2], x.shape[3])
+
+        # Calculate attention mask and relative position index in advance to speed up inference.
+        # The original code is very time-cosuming for large window size.
+        attn_mask = self.calculate_mask(x_size).to(x.device)
+        params = {
+            "attn_mask": attn_mask,
+            "rpi_sa": self.relative_position_index_SA,
+            "rpi_oca": self.relative_position_index_OCA,
+        }
+
+        x = self.patch_embed(x)
+        if self.ape is True:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        for layer in self.layers:
+            x = layer(x, x_size, params)
+
+        x = self.norm(x)  # b seq_len c
+        x = self.patch_unembed(x, x_size)
+
+        return x
+
+    def forward(self, x):
+        self.mean = self.mean.type_as(x)
+        x = (x - self.mean) * self.img_range
+
+        if self.upsampler == "pixelshuffle":
+            # for classical SR
+            x = self.conv_first(x)
+            x = self.conv_after_body(self.forward_features(x)) + x
+            x = self.conv_before_upsample(x)
+            x = self.conv_last(self.upsample(x))
+
+        x = x / self.img_range + self.mean
+
+        return x
+
+
+class ChannelAttention(nn.Module):
+    """Channel attention used in RCAN.
+
+    Args:
+    ----
+        num_feat (int): Channel number of intermediate features.
+        squeeze_factor (int): Channel squeeze factor. Default: 16.
+
+    """
+
+    def __init__(self, num_feat, squeeze_factor=16):
+        super(ChannelAttention, self).__init__()
+        self.attention = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(num_feat, num_feat // squeeze_factor, 1, padding=0),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(num_feat // squeeze_factor, num_feat, 1, padding=0),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        y = self.attention(x)
+        return x * y
+
+
+class CAB(nn.Module):
+    def __init__(self, num_feat, compress_ratio=3, squeeze_factor=30):
+        super(CAB, self).__init__()
+
+        self.cab = nn.Sequential(
+            nn.Conv2d(num_feat, num_feat // compress_ratio, 3, 1, 1),
+            nn.GELU(),
+            nn.Conv2d(num_feat // compress_ratio, num_feat, 3, 1, 1),
+            ChannelAttention(num_feat, squeeze_factor),
+        )
+
+    def forward(self, x):
+        return self.cab(x)
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """Args:
+    ----
+        x: (b, h, w, c)
+        window_size (int): window size
+
+    Returns
+    -------
+        windows: (num_windows*b, window_size, window_size, c)
+
+    """
+    b, h, w, c = x.shape
+    x = x.view(b, h // window_size, window_size, w // window_size, window_size, c)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, c)
+    return windows
+
+
+def window_reverse(windows, window_size, h, w):
+    """Args:
+    ----
+        windows: (num_windows*b, window_size, window_size, c)
+        window_size (int): Window size
+        h (int): Height of image
+        w (int): Width of image
+
+    Returns
+    -------
+        x: (b, h, w, c)
+
+    """
+    b = int(windows.shape[0] / (h * w / window_size / window_size))
+    x = windows.view(b, h // window_size, w // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(b, h, w, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    r"""Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+
+    Args:
+    ----
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+
+    """
+
+    def __init__(
+        self,
+        dim,
+        window_size,
+        num_heads,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
+        )  # 2*Wh-1 * 2*Ww-1, nH
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=0.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, rpi, mask=None):
+        """Args:
+        ----
+            x: input features with shape of (num_windows*b, n, c)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+
+        """
+        b_, n, c = x.shape
+        qkv = self.qkv(x).reshape(b_, n, 3, self.num_heads, c // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
+
+        relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(
+            self.window_size[0] * self.window_size[1],
+            self.window_size[0] * self.window_size[1],
+            -1,
+        )  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nw = mask.shape[0]
+            attn = attn.view(b_ // nw, nw, self.num_heads, n, n) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, n, n)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(b_, n, c)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class HAB(nn.Module):
+    r"""Hybrid Attention Block.
+
+    Args:
+    ----
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_resolution,
+        num_heads,
+        window_size=7,
+        shift_size=0,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        if min(self.input_resolution) <= self.window_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim,
+            window_size=to_2tuple(self.window_size),
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+
+        self.conv_scale = conv_scale
+        self.conv_block = CAB(num_feat=dim, compress_ratio=compress_ratio, squeeze_factor=squeeze_factor)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x, x_size, rpi_sa, attn_mask):
+        h, w = x_size
+        b, _, c = x.shape
+        # assert seq_len == h * w, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(b, h, w, c)
+
+        # Conv_X
+        conv_x = self.conv_block(x.permute(0, 3, 1, 2))
+        conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(b, h * w, c)
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            attn_mask = attn_mask
+        else:
+            shifted_x = x
+            attn_mask = None
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)  # nw*b, window_size, window_size, c
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, c)  # nw*b, window_size*window_size, c
+
+        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
+        attn_windows = self.attn(x_windows, rpi=rpi_sa, mask=attn_mask)
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, c)
+        shifted_x = window_reverse(attn_windows, self.window_size, h, w)  # b h' w' c
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            attn_x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            attn_x = shifted_x
+        attn_x = attn_x.view(b, h * w, c)
+
+        # FFN
+        x = shortcut + self.drop_path(attn_x) + conv_x * self.conv_scale
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+
+class PatchMerging(nn.Module):
+    r"""Patch Merging Layer.
+
+    Args:
+    ----
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x):
+        """x: b, h*w, c"""
+        h, w = self.input_resolution
+        b, seq_len, c = x.shape
+        assert seq_len == h * w, "input feature has wrong size"
+        assert h % 2 == 0 and w % 2 == 0, f"x size ({h}*{w}) are not even."
+
+        x = x.view(b, h, w, c)
+
+        x0 = x[:, 0::2, 0::2, :]  # b h/2 w/2 c
+        x1 = x[:, 1::2, 0::2, :]  # b h/2 w/2 c
+        x2 = x[:, 0::2, 1::2, :]  # b h/2 w/2 c
+        x3 = x[:, 1::2, 1::2, :]  # b h/2 w/2 c
+        x = torch.cat([x0, x1, x2, x3], -1)  # b h/2 w/2 4*c
+        x = x.view(b, -1, 4 * c)  # b h/2*w/2 4*c
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+
+class OCAB(nn.Module):
+    # overlapping cross-attention block
+
+    def __init__(
+        self,
+        dim,
+        input_resolution,
+        window_size,
+        overlap_ratio,
+        num_heads,
+        qkv_bias=True,
+        qk_scale=None,
+        mlp_ratio=2,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.window_size = window_size
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+        self.overlap_win_size = int(window_size * overlap_ratio) + window_size
+
+        self.norm1 = norm_layer(dim)
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.unfold = nn.Unfold(
+            kernel_size=(self.overlap_win_size, self.overlap_win_size),
+            stride=window_size,
+            padding=(self.overlap_win_size - window_size) // 2,
+        )
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros(
+                (window_size + self.overlap_win_size - 1) * (window_size + self.overlap_win_size - 1),
+                num_heads,
+            )
+        )  # 2*Wh-1 * 2*Ww-1, nH
+
+        trunc_normal_(self.relative_position_bias_table, std=0.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+        self.proj = nn.Linear(dim, dim)
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=nn.GELU)
+
+    def forward(self, x, x_size, rpi):
+        h, w = x_size
+        b, _, c = x.shape
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(b, h, w, c)
+
+        qkv = self.qkv(x).reshape(b, h, w, 3, c).permute(3, 0, 4, 1, 2)  # 3, b, c, h, w
+        q = qkv[0].permute(0, 2, 3, 1)  # b, h, w, c
+        kv = torch.cat((qkv[1], qkv[2]), dim=1)  # b, 2*c, h, w
+
+        # partition windows
+        q_windows = window_partition(q, self.window_size)  # nw*b, window_size, window_size, c
+        q_windows = q_windows.view(-1, self.window_size * self.window_size, c)  # nw*b, window_size*window_size, c
+
+        kv_windows = self.unfold(kv)  # b, c*w*w, nw
+        kv_windows = rearrange(
+            kv_windows,
+            "b (nc ch owh oww) nw -> nc (b nw) (owh oww) ch",
+            nc=2,
+            ch=c,
+            owh=self.overlap_win_size,
+            oww=self.overlap_win_size,
+        ).contiguous()  # 2, nw*b, ow*ow, c
+        k_windows, v_windows = kv_windows[0], kv_windows[1]  # nw*b, ow*ow, c
+
+        b_, nq, _ = q_windows.shape
+        _, n, _ = k_windows.shape
+        d = self.dim // self.num_heads
+        q = q_windows.reshape(b_, nq, self.num_heads, d).permute(0, 2, 1, 3)  # nw*b, nH, nq, d
+        k = k_windows.reshape(b_, n, self.num_heads, d).permute(0, 2, 1, 3)  # nw*b, nH, n, d
+        v = v_windows.reshape(b_, n, self.num_heads, d).permute(0, 2, 1, 3)  # nw*b, nH, n, d
+
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
+
+        relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(
+            self.window_size * self.window_size,
+            self.overlap_win_size * self.overlap_win_size,
+            -1,
+        )  # ws*ws, wse*wse, nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, ws*ws, wse*wse
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        attn = self.softmax(attn)
+        attn_windows = (attn @ v).transpose(1, 2).reshape(b_, nq, self.dim)
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, self.dim)
+        x = window_reverse(attn_windows, self.window_size, h, w)  # b h w c
+        x = x.view(b, h * w, self.dim)
+
+        x = self.proj(x) + shortcut
+
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class AttenBlocks(nn.Module):
+    """A series of attention blocks for one RHAG.
+
+    Args:
+    ----
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_resolution,
+        depth,
+        num_heads,
+        window_size,
+        compress_ratio,
+        squeeze_factor,
+        conv_scale,
+        overlap_ratio,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        norm_layer=nn.LayerNorm,
+        downsample=None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+
+        # build blocks
+        self.blocks = nn.ModuleList(
+            [
+                HAB(
+                    dim=dim,
+                    input_resolution=input_resolution,
+                    num_heads=num_heads,
+                    window_size=window_size,
+                    shift_size=0 if (i % 2 == 0) else window_size // 2,
+                    compress_ratio=compress_ratio,
+                    squeeze_factor=squeeze_factor,
+                    conv_scale=conv_scale,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop,
+                    attn_drop=attn_drop,
+                    drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                    norm_layer=norm_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+
+        # OCAB
+        self.overlap_attn = OCAB(
+            dim=dim,
+            input_resolution=input_resolution,
+            window_size=window_size,
+            overlap_ratio=overlap_ratio,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            mlp_ratio=mlp_ratio,
+            norm_layer=norm_layer,
+        )
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x, x_size, params):
+        for blk in self.blocks:
+            x = blk(x, x_size, params["rpi_sa"], params["attn_mask"])
+
+        x = self.overlap_attn(x, x_size, params["rpi_oca"])
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x
+
+
+class RHAG(nn.Module):
+    """Residual Hybrid Attention Group (RHAG).
+
+    Args:
+    ----
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        img_size: Input image size.
+        patch_size: Patch size.
+        resi_connection: The convolutional block before residual connection.
+
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_resolution,
+        depth,
+        num_heads,
+        window_size,
+        compress_ratio,
+        squeeze_factor,
+        conv_scale,
+        overlap_ratio,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        norm_layer=nn.LayerNorm,
+        downsample=None,
+        img_size=224,
+        patch_size=4,
+        resi_connection="1conv",
+    ):
+        super(RHAG, self).__init__()
+
+        self.dim = dim
+        self.input_resolution = input_resolution
+
+        self.residual_group = AttenBlocks(
+            dim=dim,
+            input_resolution=input_resolution,
+            depth=depth,
+            num_heads=num_heads,
+            window_size=window_size,
+            compress_ratio=compress_ratio,
+            squeeze_factor=squeeze_factor,
+            conv_scale=conv_scale,
+            overlap_ratio=overlap_ratio,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            drop=drop,
+            attn_drop=attn_drop,
+            drop_path=drop_path,
+            norm_layer=norm_layer,
+            downsample=downsample,
+        )
+
+        if resi_connection == "1conv":
+            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
+        elif resi_connection == "identity":
+            self.conv = nn.Identity()
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=0,
+            embed_dim=dim,
+            norm_layer=None,
+        )
+
+        self.patch_unembed = PatchUnEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=0,
+            embed_dim=dim,
+            norm_layer=None,
+        )
+
+    def forward(self, x, x_size, params):
+        return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size, params), x_size))) + x
+
+
+class PatchEmbed(nn.Module):
+    r"""Image to Patch Embedding
+
+    Args:
+    ----
+        img_size (int): Image size.  Default: 224.
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+
+    """
+
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [
+            img_size[0] // patch_size[0],
+            img_size[1] // patch_size[1],
+        ]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        x = x.flatten(2).transpose(1, 2)  # b Ph*Pw c
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+
+class PatchUnEmbed(nn.Module):
+    r"""Image to Patch Unembedding
+
+    Args:
+    ----
+        img_size (int): Image size.  Default: 224.
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+
+    """
+
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [
+            img_size[0] // patch_size[0],
+            img_size[1] // patch_size[1],
+        ]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+    def forward(self, x, x_size):
+        x = x.transpose(1, 2).contiguous().view(x.shape[0], self.embed_dim, x_size[0], x_size[1])  # b Ph*Pw c
+        return x
+
+
+class Upsample(nn.Sequential):
+    """Upsample module.
+
+    Args:
+    ----
+        scale (int): Scale factor. Supported scales: 2^n and 3.
+        num_feat (int): Channel number of intermediate features.
+
+    """
+
+    def __init__(self, scale, num_feat):
+        m = []
+        if (scale & (scale - 1)) == 0:  # scale = 2^n
+            for _ in range(int(math.log2(scale))):
+                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
+                m.append(nn.PixelShuffle(2))
+        elif scale == 3:
+            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
+            m.append(nn.PixelShuffle(3))
+        else:
+            raise ValueError(f"scale {scale} is not supported. " "Supported scales: 2^n and 3.")
+        super(Upsample, self).__init__(*m)
diff --git a/ccrestoration/config/__init__.py b/ccrestoration/config/__init__.py
index ef39a0a..151a778 100644
--- a/ccrestoration/config/__init__.py
+++ b/ccrestoration/config/__init__.py
@@ -14,3 +14,4 @@
 from ccrestoration.config.scunet_config import SCUNetConfig  # noqa
 from ccrestoration.config.dat_config import DATConfig  # noqa
 from ccrestoration.config.srcnn_config import SRCNNConfig  # noqa
+from ccrestoration.config.hat_config import HATConfig  # noqa
diff --git a/ccrestoration/config/hat_config.py b/ccrestoration/config/hat_config.py
new file mode 100644
index 0000000..cfd5dfa
--- /dev/null
+++ b/ccrestoration/config/hat_config.py
@@ -0,0 +1,334 @@
+from typing import Any, List, Optional, Tuple, Union
+
+from pydantic import field_validator
+from torch import nn
+
+from ccrestoration.config import CONFIG_REGISTRY
+from ccrestoration.type import ArchType, BaseConfig, ConfigType, ModelType
+
+
+class HATConfig(BaseConfig):
+    arch: Union[ArchType, str] = ArchType.HAT
+    model: Union[ModelType, str] = ModelType.HAT
+    scale: int = 2
+    patch_size: int = 1
+    in_chans: int = 3
+    img_size: Union[int, Tuple[int, ...]] = 64
+    img_range: float = 1.0
+    depth: Union[List[int], Tuple[int, ...]] = (6, 6, 6, 6)
+    embed_dim: int = 96
+    num_heads: Union[List[int], Tuple[int, ...]] = (6, 6, 6, 6)
+    window_size: int = 7
+    compress_ratio: int = 3
+    squeeze_factor: int = 30
+    conv_scale: float = 0.01
+    overlap_ratio: float = 0.5
+    mlp_ratio: float = 4.0
+    resi_connection: str = "1conv"
+    qkv_bias: bool = True
+    qk_scale: Optional[float] = None
+    drop_rate: float = 0.0
+    attn_drop_rate: float = 0.0
+    drop_path_rate: float = 0.1
+    ape: bool = False
+    patch_norm: bool = True
+    act_layer: Any = nn.GELU
+    norm_layer: Any = nn.LayerNorm
+    upsampler: str = "pixelshuffle"
+
+    @field_validator("scale")
+    def scale_match(cls, v: int) -> int:
+        if v not in [1, 2, 3, 4, 8]:
+            raise ValueError("scale factor must be one of [1, 2, 3, 4, 8]")
+        return v
+
+    @field_validator("upsampler")
+    def upsampler_match(cls, v: str) -> str:
+        if v not in ["pixelshuffle", "pixelshuffledirect", "nearest+conv", ""]:
+            raise ValueError("Upsampler must be one of ['pixelshuffle','pixelshuffledirect','nearest+conv', '']")
+        return v
+
+    @field_validator("resi_connection")
+    def resi_connection_match(cls, v: str) -> str:
+        if v not in ["1conv", "3conv"]:
+            raise ValueError("Residual connection must be one of ['1conv', '3conv']")
+        return v
+
+
+HATConfigs = [
+    # official models
+    # hat_s
+    HATConfig(
+        name=ConfigType.HAT_S_2x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_S_2x.pth",
+        hash="3e249a901c2aed6b82548875e21847ef6c015a40c814237a7a0abb10c69d5ddf",
+        scale=2,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=24,
+        squeeze_factor=24,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=144,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    HATConfig(
+        name=ConfigType.HAT_S_3x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_S_3x.pth",
+        hash="3733345caeb2b6cb0a7514d634b143f318e2441faf32316195a47e4cda67669e",
+        scale=3,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=24,
+        squeeze_factor=24,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=144,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    HATConfig(
+        name=ConfigType.HAT_S_4x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_S_4x.pth",
+        hash="a92f81bd2c0c1aaa371a6e4d6cac69e749fde2e36196885ee47a4a3667542c9a",
+        scale=4,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=24,
+        squeeze_factor=24,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=144,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    # hat_m
+    HATConfig(
+        name=ConfigType.HAT_2x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_2x.pth",
+        hash="35e9409849016ab2bfc0c549639ea4f1d4d5b2bdd75f856a925e124dd05670d0",
+        scale=2,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    HATConfig(
+        name=ConfigType.HAT_3x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_3x.pth",
+        hash="453c90b94d494d041c8bff2d222987af311f07b4a09b4dafbd9ab6492f15206c",
+        scale=3,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    HATConfig(
+        name=ConfigType.HAT_4x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_4x.pth",
+        hash="02dabea478aa5902a7170ad89350124e691bd89c91356f24b3267022622dc030",
+        scale=4,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    # hat_l
+    HATConfig(
+        name=ConfigType.HAT_L_ImageNet_pretrain_2x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_L_ImageNet_pretrain_2x.pth",
+        hash="2818c7ca8d72ec4cc5f31c93203d55252a662dd35cda34ce1a69661f97dcd38f",
+        scale=2,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    HATConfig(
+        name=ConfigType.HAT_L_ImageNet_pretrain_3x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_L_ImageNet_pretrain_3x.pth",
+        hash="78af181bedf1e805fd7517d3738bc5824f7ebdc477fc9f757b708bcf49ad4e3d",
+        scale=3,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    HATConfig(
+        name=ConfigType.HAT_L_ImageNet_pretrain_4x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_L_ImageNet_pretrain_4x.pth",
+        hash="5992bd38522f2b8faf11ea4bd8ee08de92465bb66892166576999afc36d60043",
+        scale=4,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    # hat_image_net
+    HATConfig(
+        name=ConfigType.HAT_ImageNet_pretrain_2x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_ImageNet_pretrain_2x.pth",
+        hash="82ebd911263bcc886fbef46b30cf97b92a932a27a3cba30163d4577afb09b9d7",
+        scale=2,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    HATConfig(
+        name=ConfigType.HAT_ImageNet_pretrain_3x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_ImageNet_pretrain_3x.pth",
+        hash="8469db903d464497f94419587054c9e83dff994edc409d563ab8f5a503767be8",
+        scale=3,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    HATConfig(
+        name=ConfigType.HAT_ImageNet_pretrain_4x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_ImageNet_pretrain_4x.pth",
+        hash="4ee053c42461187846dc0e93aa5abd34591c0725a8e044a59000e92ee215e833",
+        scale=4,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    # hat_real_gan
+    HATConfig(
+        name=ConfigType.HAT_Real_GAN_4x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_Real_GAN_4x.pth",
+        hash="f5b1e3bbbb05147ca2beefcc715279cb647d7976cbda67d62ea7e6e20d5ffcc7",
+        scale=4,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+    HATConfig(
+        name=ConfigType.HAT_Real_GAN_sharper_4x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/HAT_Real_GAN_sharper_4x.pth",
+        hash="5800b67136006eb8cab3b4ed7c8d73b6a195bb18e6cc709b674f9aa069c00271",
+        scale=4,
+        in_chans=3,
+        window_size=16,
+        compress_ratio=3,
+        squeeze_factor=30,
+        conv_scale=0.01,
+        overlap_ratio=0.5,
+        img_range=1.0,
+        depth=[6, 6, 6, 6, 6, 6],
+        embed_dim=180,
+        num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2,
+        upsampler="pixelshuffle",
+        resi_connection="1conv",
+    ),
+]
+
+for cfg in HATConfigs:
+    CONFIG_REGISTRY.register(cfg)
diff --git a/ccrestoration/model/__init__.py b/ccrestoration/model/__init__.py
index f715542..0a4b940 100644
--- a/ccrestoration/model/__init__.py
+++ b/ccrestoration/model/__init__.py
@@ -18,3 +18,4 @@
 from ccrestoration.model.scunet_model import SCUNetModel  # noqa
 from ccrestoration.model.dat_model import DATModel  # noqa
 from ccrestoration.model.srcnn_model import SRCNNModel  # noqa
+from ccrestoration.model.hat_model import HATModel  # noqa
diff --git a/ccrestoration/model/hat_model.py b/ccrestoration/model/hat_model.py
new file mode 100644
index 0000000..7c8ef6d
--- /dev/null
+++ b/ccrestoration/model/hat_model.py
@@ -0,0 +1,53 @@
+from typing import Any
+
+from ccrestoration.arch import HAT
+from ccrestoration.config import HATConfig
+from ccrestoration.model import MODEL_REGISTRY
+from ccrestoration.model.sr_base_model import SRBaseModel
+from ccrestoration.type import ModelType
+
+
+@MODEL_REGISTRY.register(name=ModelType.HAT)
+class HATModel(SRBaseModel):
+    def load_model(self) -> Any:
+        cfg: HATConfig = self.config
+        state_dict = self.get_state_dict()
+
+        if "params_ema" in state_dict:
+            state_dict = state_dict["params_ema"]
+        elif "params" in state_dict:
+            state_dict = state_dict["params"]
+        elif "model_state_dict" in state_dict:
+            # For APISR's model
+            state_dict = state_dict["model_state_dict"]
+
+        model = HAT(
+            img_size=cfg.img_size,
+            patch_size=cfg.patch_size,
+            in_chans=cfg.in_chans,
+            embed_dim=cfg.embed_dim,
+            depth=cfg.depth,
+            num_heads=cfg.num_heads,
+            window_size=cfg.window_size,
+            compress_ratio=cfg.compress_ratio,
+            squeeze_factor=cfg.squeeze_factor,
+            conv_scale=cfg.conv_scale,
+            overlap_ratio=cfg.overlap_ratio,
+            mlp_ratio=cfg.mlp_ratio,
+            qkv_bias=cfg.qkv_bias,
+            qk_scale=cfg.qk_scale,
+            drop_rate=cfg.drop_rate,
+            attn_drop_rate=cfg.attn_drop_rate,
+            drop_path_rate=cfg.drop_path_rate,
+            norm_layer=cfg.norm_layer,
+            ape=cfg.ape,
+            patch_norm=cfg.patch_norm,
+            upscale=cfg.scale,
+            img_range=cfg.img_range,
+            upsampler=cfg.upsampler,
+            resi_connection=cfg.resi_connection,
+        )
+
+        model.load_state_dict(state_dict)
+        model.eval().to(self.device)
+        return model
diff --git a/ccrestoration/type/arch.py b/ccrestoration/type/arch.py
index aa178ef..733ce8b 100644
--- a/ccrestoration/type/arch.py
+++ b/ccrestoration/type/arch.py
@@ -13,6 +13,7 @@ class ArchType(str, Enum):
     SCUNET = "SCUNET"
     DAT = "DAT"
     SRCNN = "SRCNN"
+    HAT = "HAT"
 
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 
diff --git a/ccrestoration/type/config.py b/ccrestoration/type/config.py
index cda314c..003e303 100644
--- a/ccrestoration/type/config.py
+++ b/ccrestoration/type/config.py
@@ -77,6 +77,22 @@ class ConfigType(str, Enum):
     SRCNN_3x = "SRCNN_3x.pth"
     SRCNN_4x = "SRCNN_4x.pth"
 
+    # HAT
+    HAT_S_2x = "HAT_S_2x.pth"
+    HAT_S_3x = "HAT_S_3x.pth"
+    HAT_S_4x = "HAT_S_4x.pth"
+    HAT_2x = "HAT_2x.pth"
+    HAT_3x = "HAT_3x.pth"
+    HAT_4x = "HAT_4x.pth"
+    HAT_Real_GAN_sharper_4x = "HAT_Real_GAN_sharper_4x.pth"
+    HAT_Real_GAN_4x = "HAT_Real_GAN_4x.pth"
+    HAT_ImageNet_pretrain_2x = "HAT_ImageNet_pretrain_2x.pth"
+    HAT_ImageNet_pretrain_3x = "HAT_ImageNet_pretrain_3x.pth"
+    HAT_ImageNet_pretrain_4x = "HAT_ImageNet_pretrain_4x.pth"
+    HAT_L_ImageNet_pretrain_2x = "HAT_L_ImageNet_pretrain_2x.pth"
+    HAT_L_ImageNet_pretrain_3x = "HAT_L_ImageNet_pretrain_3x.pth"
+    HAT_L_ImageNet_pretrain_4x = "HAT_L_ImageNet_pretrain_4x.pth"
+
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 
     # SpyNet
diff --git a/ccrestoration/type/model.py b/ccrestoration/type/model.py
index 387e4e4..d92cc96 100644
--- a/ccrestoration/type/model.py
+++ b/ccrestoration/type/model.py
@@ -12,6 +12,7 @@ class ModelType(str, Enum):
     SCUNet = "SCUNet"
     DAT = "DAT"
     SRCNN = "SRCNN"
+    HAT = "HAT"
 
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 
diff --git a/tests/test_hat.py b/tests/test_hat.py
new file mode 100644
index 0000000..c939728
--- /dev/null
+++ b/tests/test_hat.py
@@ -0,0 +1,76 @@
+import os
+
+import cv2
+import pytest
+
+from ccrestoration import AutoConfig, AutoModel, BaseConfig, ConfigType
+from ccrestoration.model import SRBaseModel
+
+from .util import ASSETS_PATH, calculate_image_similarity, compare_image_size, get_device, load_image
+
+
+class Test_HAT:
+    def test_official_light(self) -> None:
+        img1 = load_image()
+
+        for k in [
+            ConfigType.HAT_S_2x,
+            ConfigType.HAT_S_3x,
+            ConfigType.HAT_S_4x,
+        ]:
+            print(f"Testing {k}")
+            cfg: BaseConfig = AutoConfig.from_pretrained(k)
+            model: SRBaseModel = AutoModel.from_config(config=cfg, fp16=False, device=get_device())
+            print(model.device)
+
+            img2 = model.inference_image(img1)
+            cv2.imwrite(str(ASSETS_PATH / f"test_{k}_out.jpg"), img2)
+
+            assert calculate_image_similarity(img1, img2)
+            assert compare_image_size(img1, img2, cfg.scale)
+
+    @pytest.mark.skipif(os.environ.get("GITHUB_ACTIONS") == "true", reason="Skip on CI test")
+    def test_official(self) -> None:
+        img1 = load_image()
+
+        for k in [
+            ConfigType.HAT_S_2x,
+            ConfigType.HAT_S_3x,
+            ConfigType.HAT_S_4x,
+            ConfigType.HAT_2x,
+            ConfigType.HAT_3x,
+            ConfigType.HAT_4x,
+            ConfigType.HAT_ImageNet_pretrain_2x,
+            ConfigType.HAT_ImageNet_pretrain_3x,
+            ConfigType.HAT_ImageNet_pretrain_4x,
+            ConfigType.HAT_L_ImageNet_pretrain_2x,
+            ConfigType.HAT_L_ImageNet_pretrain_3x,
+            ConfigType.HAT_L_ImageNet_pretrain_4x,
+            ConfigType.HAT_Real_GAN_sharper_4x,
+            ConfigType.HAT_Real_GAN_4x,
+        ]:
+            print(f"Testing {k}")
+            cfg: BaseConfig = AutoConfig.from_pretrained(k)
+            model: SRBaseModel = AutoModel.from_config(config=cfg, fp16=False, device=get_device())
+            print(model.device)
+
+            img2 = model.inference_image(img1)
+            cv2.imwrite(str(ASSETS_PATH / f"test_{k}_out.jpg"), img2)
+
+            assert calculate_image_similarity(img1, img2)
+            assert compare_image_size(img1, img2, cfg.scale)
+
+    def test_custom(self) -> None:
+        img1 = load_image()
+
+        for k in [ConfigType.HAT_Real_GAN_sharper_4x]:
+            print(f"Testing {k}")
+            cfg: BaseConfig = AutoConfig.from_pretrained(k)
+            model: SRBaseModel = AutoModel.from_config(config=cfg, fp16=False, device=get_device())
+            print(model.device)
+
+            img2 = model.inference_image(img1)
+            cv2.imwrite(str(ASSETS_PATH / f"test_{k}_out.jpg"), img2)
+
+            assert calculate_image_similarity(img1, img2)
+            assert compare_image_size(img1, img2, cfg.scale)