Upload 12 files

README.md

@@ -1,13 +1,9 @@
----
-title: Low Light Deblurring New
-emoji: 📉
+title: DarkIR
+emoji: 🌻
 colorFrom: red
 colorTo: gray
 sdk: gradio
-sdk_version: 4.38.1
+sdk_version: 5.8.0
 app_file: app.py
 pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+license: mit

app.py

@@ -1,84 +1,89 @@
 import gradio as gr 
 from PIL import Image
-import os
 import torch
-import torch.nn.functional as F
 import torchvision.transforms as transforms
-import torchvision
-import numpy as np
-import yaml
-from huggingface_hub import hf_hub_download
+import torch.nn.functional as F
+
+from archs import DarkIR
 
-from archs import Network
-from options.options import parse
 
-path_opt = './options/predict/LOLBlur.yml'
 
-opt = parse(path_opt)
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 #define some auxiliary functions
 pil_to_tensor = transforms.ToTensor()
+tensor_to_pil = transforms.ToPILImage()
 
-# PATH_MODEL = opt['save']['best']
+network = 'DarkIR'
 
-model = Network(img_channel=opt['network']['img_channels'], 
-                    width=opt['network']['width'], 
-                    middle_blk_num_enc=opt['network']['middle_blk_num_enc'],
-                    middle_blk_num_dec=opt['network']['middle_blk_num_dec'],
-                    enc_blk_nums=opt['network']['enc_blk_nums'],
-                    dec_blk_nums=opt['network']['dec_blk_nums'], 
-                    dilations=opt['network']['dilations'],
-                    extra_depth_wise = opt['network']['extra_depth_wise'])
+PATH_MODEL = './models/darkir_1k_allv2_251205.pt'
 
-checkpoints = torch.load('Network_noFAC_LOLBlur.pt', map_location=device)
-# print(checkpoints)
-model.load_state_dict(checkpoints['model_state_dict'])
+model = DarkIR(img_channel=3, 
+                    width=32, 
+                    middle_blk_num_enc=2,
+                    middle_blk_num_dec=2, 
+                    enc_blk_nums=[1, 2, 3],
+                    dec_blk_nums=[3, 1, 1], 
+                    dilations=[1, 4, 9],
+                    extra_depth_wise=True)
+
+checkpoints = torch.load(PATH_MODEL, map_location=device)
+model.load_state_dict(checkpoints['params'])
 
 model = model.to(device)
 
-def load_img (filename):
-    img = Image.open(filename).convert("RGB")
-    img_tensor = pil_to_tensor(img)
-    return img_tensor
+def path_to_tensor(path):
+    img = Image.open(path).convert('RGB')
+    img = pil_to_tensor(img).unsqueeze(0)
+    
+    return img
+def normalize_tensor(tensor):
+    
+    max_value = torch.max(tensor)
+    min_value = torch.min(tensor)
+    output = (tensor - min_value)/(max_value)
+    return output
+
+def pad_tensor(tensor, multiple = 8):
+    '''pad the tensor to be multiple of some number'''
+    multiple = multiple
+    _, _, H, W = tensor.shape
+    pad_h = (multiple - H % multiple) % multiple
+    pad_w = (multiple - W % multiple) % multiple
+    tensor = F.pad(tensor, (0, pad_w, 0, pad_h), value = 0)
+    
+    return tensor
 
 def process_img(image):
-    img = np.array(image)
-    img = img / 255.
-    img = img.astype(np.float32)
-    y = torch.tensor(img).permute(2,0,1).unsqueeze(0).to(device)
+    tensor = path_to_tensor(image).to(device)
+    _, _, H, W = tensor.shape
+    
+    tensor = pad_tensor(tensor)
 
     with torch.no_grad():
-        x_hat = model(y)
-
-    restored_img = x_hat.squeeze().permute(1,2,0).clamp_(0, 1).cpu().detach().numpy()
-    restored_img = np.clip(restored_img, 0. , 1.)
+        output = model(tensor, side_loss=False)
 
-    restored_img = (restored_img * 255.0).round().astype(np.uint8)  # float32 to uint8
-    return Image.fromarray(restored_img) #(image, Image.fromarray(restored_img))
+    output = torch.clamp(output, 0., 1.)
+    output = output[:,:, :H, :W].squeeze(0)    
+    return tensor_to_pil(output)
 
-title = "Low-Light-Deblurring 🌚🌠🌝"
-description = ''' ## [Low Light Image deblurring enhancement](https://github.com/cidautai/Net-Low-light-Deblurring)
+title = "DarkIR ✏️🖼️ 🤗"
+description = ''' ## [ DarkIR: Robust Low-Light Image Restoration](https://github.com/cidautai/DarkIR)
 
 [Daniel Feijoo](https://github.com/danifei)
 
 Fundación Cidaut
 
-This model enhances low light images into normal light conditions ones. It was trained using LOLv2-real, LOLv2-synth and LOLBlur. 
-Due to the training on LOLBlur, this network is expected to also reconstruct blurred low light images. 
 
 > **Disclaimer:** please remember this is not a product, thus, you will notice some limitations.
-**This demo expects an image with some degradations.**
-Due to the CPU limitations, the model won't return results inmediately <br>.
-Except for the LOLv2-real, the model was trained using mostly synthetic data, thus it might not work great on real-world complex images. 
+**This demo expects an image with some Low-Light degradations.**
 
 <br>
 '''
 
-examples = [['examples/inputs/0010.png'],
-            ['examples/inputs/0060.png'], 
-            ['examples/inputs/0075.png'], 
-            ["examples/inputs/0087.png"], 
-            ["examples/inputs/0088.png"]]
+examples = [['examples/0010.png'],
+            ['examples/r13073518t_low.png'], 
+            ['examples/low00733_low.png'], 
+            ["examples/0087.png"]]
 
 css = """
     .image-frame img, .image-container img {

archs/DarkIR.py

@@ -0,0 +1,154 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+try:
+    from arch_model import EBlock, DBlock
+    from arch_util import CustomSequential
+except:
+    from archs.arch_model import EBlock, DBlock
+    from .arch_util import CustomSequential
+
+class DarkIR(nn.Module):
+    
+    def __init__(self, img_channel=3, 
+                 width=32, 
+                 middle_blk_num_enc=2,
+                 middle_blk_num_dec=2, 
+                 enc_blk_nums=[1, 2, 3], 
+                 dec_blk_nums=[3, 1, 1],  
+                 dilations = [1, 4, 9], 
+                 extra_depth_wise = True):
+        super(DarkIR, self).__init__()
+        
+        self.intro = nn.Conv2d(in_channels=img_channel, out_channels=width, kernel_size=3, padding=1, stride=1, groups=1,
+                                bias=True)
+        self.ending = nn.Conv2d(in_channels=width, out_channels=img_channel, kernel_size=3, padding=1, stride=1, groups=1,
+                              bias=True)
+
+        self.encoders = nn.ModuleList()
+        self.decoders = nn.ModuleList()
+        self.middle_blks = nn.ModuleList()
+        self.ups = nn.ModuleList()
+        self.downs = nn.ModuleList()
+        
+        chan = width
+        for num in enc_blk_nums:
+            self.encoders.append(
+                CustomSequential(
+                    *[EBlock(chan, extra_depth_wise=extra_depth_wise) for _ in range(num)]
+                )
+            )
+            self.downs.append(
+                nn.Conv2d(chan, 2*chan, 2, 2)
+            )
+            chan = chan * 2
+
+        self.middle_blks_enc = \
+            CustomSequential(
+                *[EBlock(chan, extra_depth_wise=extra_depth_wise) for _ in range(middle_blk_num_enc)]
+            )
+        self.middle_blks_dec = \
+            CustomSequential(
+                *[DBlock(chan, dilations=dilations, extra_depth_wise=extra_depth_wise) for _ in range(middle_blk_num_dec)]
+            )
+
+        for num in dec_blk_nums:
+            self.ups.append(
+                nn.Sequential(
+                    nn.Conv2d(chan, chan * 2, 1, bias=False),
+                    nn.PixelShuffle(2)
+                )
+            )
+            chan = chan // 2
+            self.decoders.append(
+                CustomSequential(
+                    *[DBlock(chan, dilations=dilations, extra_depth_wise=extra_depth_wise) for _ in range(num)]
+                )
+            )
+        self.padder_size = 2 ** len(self.encoders)        
+        
+        # this layer is needed for the computing of the middle loss. It isn't necessary for anything else
+        self.side_out = nn.Conv2d(in_channels = width * 2**len(self.encoders), out_channels = img_channel, 
+                                kernel_size = 3, stride=1, padding=1)
+        
+    def forward(self, input, side_loss = False, use_adapter = None):
+
+        _, _, H, W = input.shape
+
+        input = self.check_image_size(input)
+        x = self.intro(input)
+        
+        skips = []
+        for encoder, down in zip(self.encoders, self.downs):
+            x = encoder(x)
+            skips.append(x)
+            x = down(x)
+
+        # we apply the encoder transforms
+        x_light = self.middle_blks_enc(x)
+        
+        if side_loss:
+            out_side = self.side_out(x_light)
+        # apply the decoder transforms
+        x = self.middle_blks_dec(x_light)
+        x = x + x_light
+
+        for decoder, up, skip in zip(self.decoders, self.ups, skips[::-1]):
+            x = up(x)
+            x = x + skip
+            x = decoder(x)
+
+        x = self.ending(x)
+        x = x + input
+        out = x[:, :, :H, :W] # we recover the original size of the image
+        if side_loss:
+            return out_side, out
+        else:        
+            return out
+
+    def check_image_size(self, x):
+        _, _, h, w = x.size()
+        mod_pad_h = (self.padder_size - h % self.padder_size) % self.padder_size
+        mod_pad_w = (self.padder_size - w % self.padder_size) % self.padder_size
+        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), value = 0)
+        return x      
+
+if __name__ == '__main__':
+    
+    img_channel = 3
+    width = 64
+    
+    enc_blks = [1, 2, 3]
+    middle_blk_num_enc = 2
+    middle_blk_num_dec = 2
+    dec_blks = [3, 1, 1]
+    residual_layers = None
+    dilations = [1, 4, 9]
+    extra_depth_wise = True
+    
+    net = DarkIR(img_channel=img_channel, 
+                  width=width, 
+                  middle_blk_num_enc=middle_blk_num_enc,
+                  middle_blk_num_dec= middle_blk_num_dec,
+                  enc_blk_nums=enc_blks, 
+                  dec_blk_nums=dec_blks,
+                  dilations = dilations,
+                  extra_depth_wise = extra_depth_wise)
+    
+    new_state_dict = net.state_dict()
+
+    inp_shape = (3, 256, 256)
+
+    net.load_state_dict(new_state_dict)
+
+    from ptflops import get_model_complexity_info
+
+    macs, params = get_model_complexity_info(net, inp_shape, verbose=False, print_per_layer_stat=False)
+
+    print(macs, params)    
+    
+    weights = net.state_dict()
+    adapter_weights = {k: v for k, v in weights.items() if 'adapter' not in k}
+
+    
+    

archs/__init__.py

@@ -1,4 +1,232 @@
-from .nafnet_utils.arch_model import NAFNet
-from .network import Network
+import torch
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torch.nn.parallel import DistributedDataParallel as DDP
+from ptflops import get_model_complexity_info
 
-__all__ = ['NAFNet','Network']
+from .DarkIR import DarkIR   
+
+def create_model(opt, rank, adapter = False):
+    '''
+    Creates the model.
+    opt: a dictionary from the yaml config key network
+    '''
+    name = opt['name']
+
+
+    model = DarkIR(img_channel=opt['img_channels'], 
+                    width=opt['width'], 
+                    middle_blk_num_enc=opt['middle_blk_num_enc'],
+                    middle_blk_num_dec=opt['middle_blk_num_dec'], 
+                    enc_blk_nums=opt['enc_blk_nums'],
+                    dec_blk_nums=opt['dec_blk_nums'], 
+                    dilations=opt['dilations'],
+                    extra_depth_wise=opt['extra_depth_wise'])
+
+    if rank ==0:
+        print(f'Using {name} network')
+
+        input_size = (3, 256, 256)
+        macs, params = get_model_complexity_info(model, input_size, print_per_layer_stat = False)
+        print(f'Computational complexity at {input_size}: {macs}')
+        print('Number of parameters: ', params)    
+    else:
+        macs, params = None, None
+
+    model.to(rank)
+    
+    model = DDP(model, device_ids=[rank], find_unused_parameters=adapter)
+    
+    return model, macs, params
+
+def create_optim_scheduler(opt, model):
+    '''
+    Returns the optim and its scheduler.
+    opt: a dictionary of the yaml config file with the train key
+    '''
+    optim = torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()) , 
+                            lr = opt['lr_initial'],
+                            weight_decay = opt['weight_decay'],
+                            betas = opt['betas'])
+    
+    if opt['lr_scheme'] == 'CosineAnnealing':
+        scheduler = CosineAnnealingLR(optim, T_max=opt['epochs'], eta_min=opt['eta_min'])
+    else: 
+        raise NotImplementedError('scheduler not implemented')    
+        
+    return optim, scheduler
+
+def load_weights(model, old_weights):
+    '''
+    Loads the weights of a pretrained model, picking only the weights that are
+    in the new model.
+    '''
+    new_weights = model.state_dict()
+    new_weights.update({k: v for k, v in old_weights.items() if k in new_weights})
+    
+    model.load_state_dict(new_weights)
+    return model
+
+def load_optim(optim, optim_weights):
+    '''
+    Loads the values of the optimizer picking only the weights that are in the new model.
+    '''
+    optim_new_weights = optim.state_dict()
+    # optim_new_weights.load_state_dict(optim_weights)
+    optim_new_weights.update({k:v for k, v in optim_weights.items() if k in optim_new_weights})
+    return optim
+
+def resume_model(model,
+                 optim,
+                 scheduler, 
+                 path_model, 
+                 rank,resume:str=None):
+    '''
+    Returns the loaded weights of model and optimizer if resume flag is True
+    '''
+    map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
+    if resume:
+        checkpoints = torch.load(path_model, map_location=map_location, weights_only=False)
+        weights = checkpoints['model_state_dict']
+        model = load_weights(model, old_weights=weights)
+        optim = load_optim(optim, optim_weights = checkpoints['optimizer_state_dict'])
+        scheduler.load_state_dict(checkpoints['scheduler_state_dict'])
+        start_epochs = checkpoints['epoch']
+
+        if rank == 0: print('Loaded weights')
+    else:
+        start_epochs = 0
+        if rank==0: print('Starting from zero the training')
+    
+    return model, optim, scheduler, start_epochs
+
+def find_different_keys(dict1, dict2):
+
+# Finding different keys
+    different_keys = set(dict1.keys()) ^ set(dict2.keys())
+
+    return different_keys
+
+def number_common_keys(dict1, dict2):
+    # Finding common keys
+    common_keys = set(dict1.keys()) & set(dict2.keys())
+
+    # Counting the number of common keys
+    common_keys_count = len(common_keys)
+    return common_keys_count
+
+# # Function to add 'modules_list' prefix after the first numeric index
+# def add_middle_prefix(state_dict, middle_prefix, target_strings):
+#     new_state_dict = {}
+#     for key, value in state_dict.items():
+#         for target in target_strings:
+#             if target in key:
+#                 parts = key.split('.')
+#                 # Find the first numeric index after the target string
+#                 for i, part in enumerate(parts):
+#                     if part == target:
+#                         # Insert the middle prefix after the first numeric index
+#                         if i + 1 < len(parts) and parts[i + 1].isdigit():
+#                             parts.insert(i + 2, middle_prefix)
+#                             break
+#                 new_key = '.'.join(parts)
+#                 new_state_dict[new_key] = value
+#                 break
+#         else:
+#             new_state_dict[key] = value
+#     return new_state_dict
+
+# # Function to adjust keys for 'middle_blks.' prefix
+# def adjust_middle_blks_keys(state_dict, target_prefix, middle_prefix):
+#     new_state_dict = {}
+#     for key, value in state_dict.items():
+#         if target_prefix in key:
+#             parts = key.split('.')
+#             # Find the target prefix and adjust the key
+#             for i, part in enumerate(parts):
+#                 if part == target_prefix.rstrip('.'):
+#                     if i + 1 < len(parts) and parts[i + 1].isdigit():
+#                         # Swap the numerical part and the middle prefix
+#                         new_key = '.'.join(parts[:i + 1] + [middle_prefix] + parts[i + 1:i + 2] + parts[i + 2:])
+#                         new_state_dict[new_key] = value
+#                         break
+#         else:
+#             new_state_dict[key] = value
+#     return new_state_dict
+
+# def resume_nafnet(model,
+#                  optim,
+#                  scheduler, 
+#                  path_adapter,
+#                  path_model,
+#                  rank, resume:str=None):
+#     '''
+#     Returns the loaded weights of model and optimizer if resume flag is True
+#     '''
+#     map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
+#     #first load the model weights
+#     checkpoints = torch.load(path_model, map_location=map_location, weights_only=False)
+#     weights = checkpoints
+#     if rank==0:
+#         print(len(weights), len(model.state_dict().keys()))
+
+#         different_keys = find_different_keys(weights, model.state_dict())
+#         filtered_keys = {item for item in different_keys if 'adapter' not in item}
+#         print(filtered_keys)
+#         print(len(filtered_keys))
+#     model = load_weights(model, old_weights=weights) 
+#     #now if needed load the adapter weights
+#     if resume:
+#         checkpoints = torch.load(path_adapter, map_location=map_location, weights_only=False)
+#         weights = checkpoints
+#         model = load_weights(model, old_weights=weights)
+#         # optim = load_optim(optim, optim_weights = checkpoints['optimizer_state_dict'])
+#         scheduler.load_state_dict(checkpoints['scheduler_state_dict'])
+#         start_epochs = checkpoints['epoch']
+
+#         if rank == 0: print('Loaded weights')
+#     else:
+#         start_epochs = 0
+#         if rank == 0: print('Starting from zero the training')
+    
+#     return model, optim, scheduler, start_epochs
+
+def save_checkpoint(model, optim, scheduler, metrics_eval, metrics_train, paths, adapter = False, rank = None):
+
+    '''
+    Save the .pt of the model after each epoch.
+    '''
+    best_psnr = metrics_train['best_psnr']
+    if rank!=0: 
+        return best_psnr
+    
+    if type(next(iter(metrics_eval.values()))) != dict:
+        metrics_eval = {'metrics': metrics_eval}
+
+    weights = model.state_dict()
+
+    # Save the model after every epoch
+    model_to_save = {
+        'epoch': metrics_train['epoch'],
+        'model_state_dict': weights,
+        'optimizer_state_dict': optim.state_dict(),
+        'loss': metrics_train['train_loss'],
+        'scheduler_state_dict': scheduler.state_dict()
+    }
+
+    try:
+        torch.save(model_to_save, paths['new'])
+
+        # Save best model if new valid_psnr is higher than the best one
+        if next(iter(metrics_eval.values()))['valid_psnr'] >= metrics_train['best_psnr']:
+            torch.save(model_to_save, paths['best'])
+            metrics_train['best_psnr'] = next(iter(metrics_eval.values()))['valid_psnr']  # update best psnr
+    except Exception as e:
+        print(f"Error saving model: {e}")
+    return metrics_train['best_psnr']
+
+__all__ = ['create_model', 'resume_model', 'create_optim_scheduler', 'save_checkpoint',
+           'load_optim', 'load_weights']
+
+
+
+    

archs/arch_model.py

@@ -0,0 +1,240 @@
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import torch.nn.functional as F
+
+try:
+    from .arch_util import LayerNorm2d
+except:
+    from arch_util import LayerNorm2d
+
+
+class SimpleGate(nn.Module):
+    def forward(self, x):
+        x1, x2 = x.chunk(2, dim=1)
+        return x1 * x2
+
+class Adapter(nn.Module):
+    
+    def __init__(self, c, ffn_channel = None):
+        super().__init__()
+        if ffn_channel:
+            ffn_channel = 2
+        else:
+            ffn_channel = c
+        self.conv1 = nn.Conv2d(in_channels=c, out_channels=ffn_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
+        self.conv2 = nn.Conv2d(in_channels=ffn_channel, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
+        self.depthwise = nn.Conv2d(in_channels=c, out_channels=ffn_channel, kernel_size=3, padding=1, stride=1, groups=c, bias=True, dilation=1)
+
+    def forward(self, input):
+        
+        x = self.conv1(input) + self.depthwise(input)
+        x = self.conv2(x)
+        
+        return x
+
+class FreMLP(nn.Module):
+    
+    def __init__(self, nc, expand = 2):
+        super(FreMLP, self).__init__()
+        self.process1 = nn.Sequential(
+            nn.Conv2d(nc, expand * nc, 1, 1, 0),
+            nn.LeakyReLU(0.1, inplace=True),
+            nn.Conv2d(expand * nc, nc, 1, 1, 0))
+
+    def forward(self, x):
+        _, _, H, W = x.shape
+        x_freq = torch.fft.rfft2(x, norm='backward')
+        mag = torch.abs(x_freq)
+        pha = torch.angle(x_freq)
+        mag = self.process1(mag)
+        real = mag * torch.cos(pha)
+        imag = mag * torch.sin(pha)
+        x_out = torch.complex(real, imag)
+        x_out = torch.fft.irfft2(x_out, s=(H, W), norm='backward')
+        return x_out
+
+class Branch(nn.Module):
+    '''
+    Branch that lasts lonly the dilated convolutions
+    '''
+    def __init__(self, c, DW_Expand, dilation = 1):
+        super().__init__()
+        self.dw_channel = DW_Expand * c 
+        
+        self.branch = nn.Sequential(
+                       nn.Conv2d(in_channels=self.dw_channel, out_channels=self.dw_channel, kernel_size=3, padding=dilation, stride=1, groups=self.dw_channel,
+                                            bias=True, dilation = dilation) # the dconv
+        )
+    def forward(self, input):
+        return self.branch(input)
+    
+class DBlock(nn.Module):
+    '''
+    Change this block using Branch
+    '''
+    
+    def __init__(self, c, DW_Expand=2, FFN_Expand=2, dilations = [1], extra_depth_wise = False):
+        super().__init__()
+        #we define the 2 branches
+        self.dw_channel = DW_Expand * c 
+
+        self.conv1 = nn.Conv2d(in_channels=c, out_channels=self.dw_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True, dilation = 1)
+        self.extra_conv = nn.Conv2d(self.dw_channel, self.dw_channel, kernel_size=3, padding=1, stride=1, groups=c, bias=True, dilation=1) if extra_depth_wise else nn.Identity() #optional extra dw
+        self.branches = nn.ModuleList()
+        for dilation in dilations:
+            self.branches.append(Branch(self.dw_channel, DW_Expand = 1, dilation = dilation))
+            
+        assert len(dilations) == len(self.branches)
+        self.dw_channel = DW_Expand * c 
+        self.sca = nn.Sequential(
+                       nn.AdaptiveAvgPool2d(1),
+                       nn.Conv2d(in_channels=self.dw_channel // 2, out_channels=self.dw_channel // 2, kernel_size=1, padding=0, stride=1,
+                       groups=1, bias=True, dilation = 1),  
+        )
+        self.sg1 = SimpleGate()
+        self.sg2 = SimpleGate()
+        self.conv3 = nn.Conv2d(in_channels=self.dw_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True, dilation = 1)
+        ffn_channel = FFN_Expand * c
+        self.conv4 = nn.Conv2d(in_channels=c, out_channels=ffn_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
+        self.conv5 = nn.Conv2d(in_channels=ffn_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
+
+        self.norm1 = LayerNorm2d(c)
+        self.norm2 = LayerNorm2d(c)
+
+        self.gamma = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
+        self.beta = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
+
+
+#        self.adapter = Adapter(c, ffn_channel=None)
+        
+#        self.use_adapters = False
+
+#    def set_use_adapters(self, use_adapters):
+#        self.use_adapters = use_adapters
+        
+    def forward(self, inp, adapter = None):
+
+        y = inp
+        x = self.norm1(inp)
+        # x = self.conv1(self.extra_conv(x))
+        x = self.extra_conv(self.conv1(x))
+        z = 0
+        for branch in self.branches:
+            z += branch(x)
+        
+        z = self.sg1(z)
+        x = self.sca(z) * z
+        x = self.conv3(x)
+        y = inp + self.beta * x
+        #second step
+        x = self.conv4(self.norm2(y)) # size [B, 2*C, H, W]
+        x = self.sg2(x)  # size [B, C, H, W]
+        x = self.conv5(x) # size [B, C, H, W]
+        x = y + x * self.gamma
+        
+#        if self.use_adapters:
+#            return self.adapter(x)
+#        else:
+        return x 
+
+class EBlock(nn.Module):
+    '''
+    Change this block using Branch
+    '''
+    
+    def __init__(self, c, DW_Expand=2, dilations = [1], extra_depth_wise = False):
+        super().__init__()
+        #we define the 2 branches
+        self.dw_channel = DW_Expand * c 
+        self.extra_conv = nn.Conv2d(c, c, kernel_size=3, padding=1, stride=1, groups=c, bias=True, dilation=1) if extra_depth_wise else nn.Identity() #optional extra dw
+        self.conv1 = nn.Conv2d(in_channels=c, out_channels=self.dw_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True, dilation = 1)
+                
+        self.branches = nn.ModuleList()
+        for dilation in dilations:
+            self.branches.append(Branch(c, DW_Expand, dilation = dilation))
+            
+        assert len(dilations) == len(self.branches)
+        self.dw_channel = DW_Expand * c 
+        self.sca = nn.Sequential(
+                       nn.AdaptiveAvgPool2d(1),
+                       nn.Conv2d(in_channels=self.dw_channel // 2, out_channels=self.dw_channel // 2, kernel_size=1, padding=0, stride=1,
+                       groups=1, bias=True, dilation = 1),  
+        )
+        self.sg1 = SimpleGate()
+        self.conv3 = nn.Conv2d(in_channels=self.dw_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True, dilation = 1)
+        # second step
+
+        self.norm1 = LayerNorm2d(c)
+        self.norm2 = LayerNorm2d(c)
+        self.freq = FreMLP(nc = c, expand=2)
+        self.gamma = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
+        self.beta = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
+
+
+#        self.adapter = Adapter(c, ffn_channel=None)
+        
+#        self.use_adapters = False
+
+#    def set_use_adapters(self, use_adapters):
+#        self.use_adapters = use_adapters
+
+    def forward(self, inp):
+        y = inp
+        x = self.norm1(inp)
+        x = self.conv1(self.extra_conv(x))
+        z = 0
+        for branch in self.branches:
+            z += branch(x)
+        
+        z = self.sg1(z)
+        x = self.sca(z) * z
+        x = self.conv3(x)
+        y = inp + self.beta * x
+        #second step
+        x_step2 = self.norm2(y) # size [B, 2*C, H, W]
+        x_freq = self.freq(x_step2) # size [B, C, H, W]
+        x = y * x_freq 
+        x = y + x * self.gamma
+
+#        if self.use_adapters:
+#            return self.adapter(x)
+#        else:
+        return x 
+
+#----------------------------------------------------------------------------------------------
+if __name__ == '__main__':
+    
+    img_channel = 3
+    width = 32
+
+    enc_blks = [1, 2, 3]
+    middle_blk_num = 3
+    dec_blks = [3, 1, 1]
+    dilations = [1, 4, 9]
+    extra_depth_wise = True
+    
+    # net = NAFNet(img_channel=img_channel, width=width, middle_blk_num=middle_blk_num,
+    #                   enc_blk_nums=enc_blks, dec_blk_nums=dec_blks)
+    net  = EBlock(c = img_channel, 
+                            dilations = dilations,
+                            extra_depth_wise=extra_depth_wise)
+
+    inp_shape = (3, 256, 256)
+
+    from ptflops import get_model_complexity_info
+
+    macs, params = get_model_complexity_info(net, inp_shape, verbose=False, print_per_layer_stat=False)
+    output = net(torch.randn((4, 3, 256, 256)))
+    # print('Values of EBlock:')
+    print(macs, params)
+
+    channels = 128
+    resol = 32
+    ksize = 5
+
+    # net = FAC(channels=channels, ksize=ksize)
+    # inp_shape = (channels, resol, resol)
+    # macs, params = get_model_complexity_info(net, inp_shape, verbose=False, print_per_layer_stat=True)
+    # print('Values of FAC:')
+    # print(macs, params)

archs/arch_util.py

@@ -1,118 +1,65 @@
 import torch
-import torch.nn as nn
-import torch.nn.init as init
-import torch.nn.functional as F
-
-try:
-    from .nafnet_utils.arch_util import LayerNorm2d
-    from .nafnet_utils.arch_model import SimpleGate
-except:
-    from nafnet_utils.arch_util import LayerNorm2d
-    from nafnet_utils.arch_model import SimpleGate
-
-
-class Branch(nn.Module):
-    '''
-    Branch that lasts lonly the dilated convolutions
-    '''
-    def __init__(self, c, DW_Expand, dilation = 1, extra_depth_wise = False):
-        super().__init__()
-        self.dw_channel = DW_Expand * c 
-        self.branch = nn.Sequential(
-                       nn.Conv2d(c, c, kernel_size=3, padding=1, stride=1, groups=c, bias=True, dilation=1) if extra_depth_wise else nn.Identity(), #optional extra dw
-                       nn.Conv2d(in_channels=c, out_channels=self.dw_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True, dilation = 1),
-                       nn.Conv2d(in_channels=self.dw_channel, out_channels=self.dw_channel, kernel_size=3, padding=dilation, stride=1, groups=self.dw_channel,
-                                            bias=True, dilation = dilation) # the dconv
-        )
-    def forward(self, input):
-        return self.branch(input)
-    
-class EBlock(nn.Module):
+import numpy as np
+from torch import nn as nn
+from torch.nn import init as init
+import torch.distributed as dist
+from collections import OrderedDict
+
+class LayerNormFunction(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, weight, bias, eps):
+        ctx.eps = eps
+        N, C, H, W = x.size()
+        mu = x.mean(1, keepdim=True)
+        var = (x - mu).pow(2).mean(1, keepdim=True)
+        y = (x - mu) / (var + eps).sqrt()
+        ctx.save_for_backward(y, var, weight)
+        y = weight.view(1, C, 1, 1) * y + bias.view(1, C, 1, 1)
+        return y
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        eps = ctx.eps
+
+        N, C, H, W = grad_output.size()
+        y, var, weight = ctx.saved_variables
+        g = grad_output * weight.view(1, C, 1, 1)
+        mean_g = g.mean(dim=1, keepdim=True)
+
+        mean_gy = (g * y).mean(dim=1, keepdim=True)
+        gx = 1. / torch.sqrt(var + eps) * (g - y * mean_gy - mean_g)
+        return gx, (grad_output * y).sum(dim=3).sum(dim=2).sum(dim=0), grad_output.sum(dim=3).sum(dim=2).sum(
+            dim=0), None
+
+class LayerNorm2d(nn.Module):
+
+    def __init__(self, channels, eps=1e-6):
+        super(LayerNorm2d, self).__init__()
+        self.register_parameter('weight', nn.Parameter(torch.ones(channels)))
+        self.register_parameter('bias', nn.Parameter(torch.zeros(channels)))
+        self.eps = eps
+
+    def forward(self, x):
+        return LayerNormFunction.apply(x, self.weight, self.bias, self.eps)
+
+
+class CustomSequential(nn.Module):
     '''
-    Change this block using Branch
+    Similar to nn.Sequential, but it lets us introduce a second argument in the forward method 
+    so adaptors can be considered in the inference.
     '''
-    
-    def __init__(self, c, DW_Expand=2, FFN_Expand=2, dilations = [1], extra_depth_wise = False):
-        super().__init__()
-        #we define the 2 branches
-        
-        self.branches = nn.ModuleList()
-        for dilation in dilations:
-            self.branches.append(Branch(c, DW_Expand, dilation = dilation, extra_depth_wise=extra_depth_wise))
-            
-        assert len(dilations) == len(self.branches)
-        self.dw_channel = DW_Expand * c 
-        self.sca = nn.Sequential(
-                       nn.AdaptiveAvgPool2d(1),
-                       nn.Conv2d(in_channels=self.dw_channel // 2, out_channels=self.dw_channel // 2, kernel_size=1, padding=0, stride=1,
-                       groups=1, bias=True, dilation = 1),  
-        )
-        self.sg1 = SimpleGate()
-        self.sg2 = SimpleGate()
-        self.conv3 = nn.Conv2d(in_channels=self.dw_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True, dilation = 1)
-        ffn_channel = FFN_Expand * c
-        self.conv4 = nn.Conv2d(in_channels=c, out_channels=ffn_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
-        self.conv5 = nn.Conv2d(in_channels=ffn_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
-
-        self.norm1 = LayerNorm2d(c)
-        self.norm2 = LayerNorm2d(c)
-
-        self.gamma = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
-        self.beta = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
-
-    def forward(self, inp):
+    def __init__(self, *args):
+        super(CustomSequential, self).__init__()
+        self.modules_list = nn.ModuleList(args)
 
-        y = inp
-        x = self.norm1(inp)
-        z = 0
-        for branch in self.branches:
-            z += branch(x)
-        
-        z = self.sg1(z)
-        x = self.sca(z) * z
-        x = self.conv3(x)
-        y = inp + self.beta * x
-        #second step
-        x = self.conv4(self.norm2(y)) # size [B, 2*C, H, W]
-        x = self.sg2(x)  # size [B, C, H, W]
-        x = self.conv5(x) # size [B, C, H, W]
+    def forward(self, x, use_adapter=False):
+        for module in self.modules_list:
+            if hasattr(module, 'set_use_adapters'):
+                module.set_use_adapters(use_adapter)
+            x = module(x)
+        return x
 
-        return y + x * self.gamma
-
-#----------------------------------------------------------------------------------------------
 if __name__ == '__main__':
     
-    img_channel = 3
-    width = 32
-
-    enc_blks = [1, 2, 3]
-    middle_blk_num = 3
-    dec_blks = [3, 1, 1]
-    dilations = [1, 4, 9]
-    extra_depth_wise = False
-    
-    # net = NAFNet(img_channel=img_channel, width=width, middle_blk_num=middle_blk_num,
-    #                   enc_blk_nums=enc_blks, dec_blk_nums=dec_blks)
-    net  = EBlock(c = img_channel, 
-                            dilations = dilations,
-                            extra_depth_wise=extra_depth_wise)
-
-    inp_shape = (3, 256, 256)
-
-    from ptflops import get_model_complexity_info
-
-    # macs, params = get_model_complexity_info(net, inp_shape, verbose=False, print_per_layer_stat=True)
-
-    # print('Values of EBlock:')
-    # print(macs, params)
-
-    channels = 128
-    resol = 32
-    ksize = 5
-
-    net = FAC(channels=channels, ksize=ksize)
-    inp_shape = (channels, resol, resol)
-    macs, params = get_model_complexity_info(net, inp_shape, verbose=False, print_per_layer_stat=True)
-    print('Values of FAC:')
-    print(macs, params)
-
+    pass

darkir_v2real+lsrw.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e506a3ef1e60d7223a3d9846ebf4851a86f34f02a927fbcbd9f97edef62a5a6
+size 13421785

requirements.txt

@@ -1,108 +1,19 @@
-aiofiles==23.2.1
-altair==5.3.0
-annotated-types==0.7.0
-anyio==4.4.0
-attrs==23.2.0
-certifi==2024.6.2
-charset-normalizer==3.3.2
-click==8.1.7
-contourpy==1.2.1
-cycler==0.12.1
-dnspython==2.6.1
-docker-pycreds==0.4.0
-email_validator==2.2.0
-exceptiongroup==1.2.1
-fastapi==0.111.0
-fastapi-cli==0.0.4
-ffmpy==0.3.2
-filelock==3.15.3
-fonttools==4.53.0
-fsspec==2024.6.0
-gitdb==4.0.11
-GitPython==3.1.43
-gradio==4.36.1
-gradio_client==1.0.1
-h11==0.14.0
-httpcore==1.0.5
-httptools==0.6.1
-httpx==0.27.0
-huggingface-hub==0.23.4
-idna==3.7
-importlib_resources==6.4.0
-Jinja2==3.1.4
-jsonschema==4.22.0
-jsonschema-specifications==2023.12.1
-kiwisolver==1.4.5
+einops==0.8.0
+gradio==5.8.0
 kornia==0.7.2
-kornia_rs==0.1.3
 lpips==0.1.4
-markdown-it-py==3.0.0
-MarkupSafe==2.1.5
-matplotlib==3.9.0
-mdurl==0.1.2
-mpmath==1.3.0
-networkx==3.3
 numpy==2.0.0
-nvidia-cublas-cu12==12.1.3.1
-nvidia-cuda-cupti-cu12==12.1.105
-nvidia-cuda-nvrtc-cu12==12.1.105
-nvidia-cuda-runtime-cu12==12.1.105
-nvidia-cudnn-cu12==8.9.2.26
-nvidia-cufft-cu12==11.0.2.54
-nvidia-curand-cu12==10.3.2.106
-nvidia-cusolver-cu12==11.4.5.107
-nvidia-cusparse-cu12==12.1.0.106
-nvidia-nccl-cu12==2.20.5
-nvidia-nvjitlink-cu12==12.5.40
-nvidia-nvtx-cu12==12.1.105
 opencv-python==4.10.0.84
-orjson==3.10.5
-packaging==24.1
 pandas==2.2.2
 pillow==10.3.0
-platformdirs==4.2.2
-protobuf==5.27.1
-psutil==6.0.0
 ptflops==0.7.3
-pydantic==2.7.4
-pydantic_core==2.18.4
-pydub==0.25.1
-Pygments==2.18.0
-pyparsing==3.1.2
-python-dateutil==2.9.0.post0
-python-dotenv==1.0.1
-python-multipart==0.0.9
+pyiqa==0.1.13
 pytorch-msssim==1.0.0
-pytz==2024.1
 PyYAML==6.0.1
-referencing==0.35.1
-requests==2.32.3
-rich==13.7.1
-rpds-py==0.18.1
-ruff==0.4.10
+scikit-image==0.24.0
 scipy==1.13.1
-semantic-version==2.10.0
-sentry-sdk==2.6.0
-setproctitle==1.3.3
-shellingham==1.5.4
-six==1.16.0
-smmap==5.0.1
-sniffio==1.3.1
-starlette==0.37.2
-sympy==1.12.1
-tomlkit==0.12.0
-toolz==0.12.1
-torch==2.3.1
-torchvision==0.18.1
+torch==2.5.1
+torchaudio==2.5.1
+torchvision==0.20.1
 tqdm==4.66.4
-triton==2.3.1
-typer==0.12.3
-typing_extensions==4.12.2
-tzdata==2024.1
-ujson==5.10.0
-urllib3==2.2.2
-uvicorn==0.30.1
-uvloop==0.19.0
-wandb==0.17.2
-watchfiles==0.22.0
-websockets==11.0.3
+wandb==0.17.2