| import torch.nn as nn | |
| from basicsr.archs.mair_arch import MaIR | |
| def buildMaIR_Small(upscale=2): | |
| return MaIR(img_size=64, | |
| patch_size=1, | |
| in_chans=3, | |
| embed_dim=60, | |
| depths=(6, 6, 6, 6), | |
| mlp_ratio=1.6, | |
| ssm_ratio=1.4, | |
| drop_rate=0., | |
| norm_layer=nn.LayerNorm, | |
| patch_norm=True, | |
| use_checkpoint=False, | |
| upscale=upscale, | |
| img_range=1., | |
| upsampler='pixelshuffledirect', | |
| resi_connection='1conv') | |
| def buildMaIR_Tiny(upscale=2): | |
| return MaIR(img_size=64, | |
| patch_size=1, | |
| in_chans=3, | |
| embed_dim=60, | |
| depths=(6, 6, 6, 6), | |
| mlp_ratio=1.6, | |
| ssm_ratio=1.1, | |
| d_state=1, | |
| drop_rate=0., | |
| norm_layer=nn.LayerNorm, | |
| patch_norm=True, | |
| use_checkpoint=False, | |
| upscale=upscale, | |
| img_range=1., | |
| upsampler='pixelshuffledirect', | |
| resi_connection='1conv') | |
| def buildMaIR(upscale=2): | |
| return MaIR(img_size=64, | |
| patch_size=1, | |
| in_chans=3, | |
| embed_dim=180, | |
| depths=(6, 6, 6, 6, 6, 6), | |
| mlp_ratio=2., | |
| drop_rate=0., | |
| norm_layer=nn.LayerNorm, | |
| patch_norm=True, | |
| use_checkpoint=False, | |
| upscale=upscale, | |
| img_range=1., | |
| upsampler='pixelshuffle', | |
| resi_connection='1conv') | |
| def buildMaIR_SR(upscale=2): | |
| return MaIR(img_size=64, | |
| patch_size=1, | |
| in_chans=3, | |
| embed_dim=180, | |
| depths=(6, 6, 6, 6, 6, 6), | |
| drop_rate=0., | |
| d_state=16, | |
| ssm_ratio=2.0, | |
| mlp_ratio=2.5, | |
| drop_path_rate=0.1, | |
| norm_layer=nn.LayerNorm, | |
| patch_norm=True, | |
| use_checkpoint=False, | |
| upscale=upscale, | |
| img_range=1., | |
| upsampler='pixelshuffle', | |
| resi_connection='1conv', | |
| dynamic_ids=False, | |
| scan_len=4, | |
| batch_size=1, | |
| ) | |