run, but doesnt seem to work that well

README.md

@@ -7,6 +7,7 @@ Click name to jump to workflow
 2. [**RealViformer**](#realviformer) - Paper: [Investigating Attention for Real-World Video Super-Resolution](https://github.com/Yuehan717/RealViformer)
 3. [**ProPIH**](#propih). Paper: [Progressive Painterly Image Harmonization from Low-level Styles to High-level Styles](https://github.com/bcmi/ProPIH-Painterly-Image-Harmonization)
 4. [**CoLIE**](#colie). Paper: [Fast Context-Based Low-Light Image Enhancement via Neural Implicit Representations](https://github.com/ctom2/colie)
+5. [**ST-AVSR**](#st-avsr). Paper: [Arbitrary-Scale Video Super-Resolution with Structural and Textural Priors](https://github.com/shangwei5/ST-AVSR)
    
 ## Workflows
 
@@ -52,3 +53,13 @@ No model needed to be downloaded. Lower `loss_mean` seems to result in brighter
 [workflow_colie_lowlight.json](workflow/workflow_colie_lowlight.json)
 
 ![colie_lowlight](workflow/colie_lowlight.png)
+
+### ST-AVSR
+
+You will need to install [mmcv](https://github.com/open-mmlab/mmcv), which can be a real pain. Refer to https://mmcv.readthedocs.io/en/latest/get_started/installation.html. Cupy is also required. This is tested on:
+
+```
+mmcv==2.2.0
+torch=2.3.0+cu121
+cupy==13.2.0
+```

__init__.py

@@ -9,6 +9,7 @@ NODE_CLASS_MAPPINGS = {
     "CoLIE_LowLight_Enhance": CoLIE_LowLight_Enhance,
 }
 
+
 NODE_DISPLAY_NAME_MAPPINGS = {
     "WFEN": "WFEN Face Restore",
     "RealViFormerSR": "RealViFormer Video SR",
@@ -16,5 +17,12 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CoLIE_LowLight_Enhance": "CoLIE LowLight Enhance",
 }
 
+try:
+    from .asvr_run import ASVR_VideoSR
+    NODE_CLASS_MAPPINGS["ASVR_VideoSR"] = ASVR_VideoSR
+    NODE_DISPLAY_NAME_MAPPINGS["ASVR_VideoSR"] = "ASVR Video SR"
+except ImportError as e:
+    print(f"Failed to load ASVR_VideoSR. This will not affect other nodes. {e}")
+    pass
 
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

asvr_module/ADISTS_allstage.py

@@ -0,0 +1,191 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from torchvision import models
+import torch.nn.functional as F
+
+class Downsample(nn.Module):
+    def __init__(self, filter_size=5, stride=2, channels=None, pad_off=0):
+        super(Downsample, self).__init__()
+        self.padding = (filter_size - 2) // 2
+        self.stride = stride
+        self.channels = channels
+        a = np.hanning(filter_size)[1:-1]
+        g = torch.Tensor(a[:, None] * a[None, :])
+        g = g / torch.sum(g)
+        self.register_buffer('filter', g[None, None, :, :].repeat((self.channels, 1, 1, 1)))
+
+    def forward(self, input):
+        input = input ** 2
+        out = F.conv2d(input, self.filter, stride=self.stride, padding=self.padding, groups=input.shape[1])
+        return (out + 1e-12).sqrt()
+
+
+class ADISTS(torch.nn.Module):
+    def __init__(self, window_size=21):
+        super(ADISTS, self).__init__()
+        vgg_pretrained_features = models.vgg16(pretrained=True).features
+        self.stage1 = torch.nn.Sequential()
+        self.stage2 = torch.nn.Sequential()
+        self.stage3 = torch.nn.Sequential()
+        self.stage4 = torch.nn.Sequential()
+        self.stage5 = torch.nn.Sequential()
+        for x in range(0, 4):
+            self.stage1.add_module(str(x), vgg_pretrained_features[x])
+        self.stage2.add_module(str(4), Downsample(channels=64))
+        for x in range(5, 9):
+            self.stage2.add_module(str(x), vgg_pretrained_features[x])
+        self.stage3.add_module(str(9), Downsample(channels=128))
+        for x in range(10, 16):
+            self.stage3.add_module(str(x), vgg_pretrained_features[x])
+        self.stage4.add_module(str(16), Downsample(channels=256))
+        for x in range(17, 23):
+            self.stage4.add_module(str(x), vgg_pretrained_features[x])
+        self.stage5.add_module(str(23), Downsample(channels=512))
+        for x in range(24, 30):
+            self.stage5.add_module(str(x), vgg_pretrained_features[x])
+
+        for param in self.parameters():
+            param.requires_grad = False
+
+        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, -1, 1, 1))
+        self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1, -1, 1, 1))
+
+        # self.chns = [3, 64]  #[3, 64, 128, 256, 512, 512]
+        # self.windows = nn.ParameterList()
+        # self.window_size = window_size
+        # for k in range(len(self.chns)):
+        #     self.windows.append(self.create_window(self.window_size, self.window_size / 3, self.chns[k]))
+
+    def forward_once(self, x):
+        h_list = []
+        h = (x - self.mean) / self.std
+        h = self.stage1(h)
+        _, _, H, W = h.size()
+        h_list.append(h)
+        h = self.stage2(h)
+        h_list.append(F.interpolate(h, (H, W), mode="bilinear"))
+        h = self.stage3(h)
+        h_list.append(F.interpolate(h, (H, W), mode="bilinear"))
+        h = self.stage4(h)
+        h_list.append(F.interpolate(h, (H, W), mode="bilinear"))
+        h = self.stage5(h)
+        h_list.append(F.interpolate(h, (H, W), mode="bilinear"))
+
+        return torch.cat(h_list, dim=1) #B, 64+12+256+512+512, H, W
+
+    def std_(self, img, window_size=3):
+        assert window_size % 2 == 1
+        pad = window_size // 2
+
+        # calculate std on the mean image of the color channels
+        # img = torch.mean(img, dim=1, keepdim=True)
+        N, C, H, W = img.shape
+        img = nn.functional.pad(img, [pad] * 4, mode='reflect')
+        img = nn.functional.unfold(img, kernel_size=window_size)
+        img = img.view(N, C, window_size * window_size, H, W)
+        img = img - torch.mean(img, dim=2, keepdim=True)
+        # print(img.shape)
+        img = img * img
+        img = torch.mean(img, dim=2, keepdim=True)
+        # print(img.shape)
+        img = torch.sqrt(img)
+        # print(img.shape)
+        img = img.squeeze(2)
+        return img
+
+    def generate_alpha(self,input, lower=1, upper=5):
+        N, C, H, W = input.shape
+        ratio = input.new_ones((N, C, H, W)) * 0.5
+        input_std = self.std_(input)
+        ratio[input_std < lower] = torch.sigmoid((input_std - lower))[input_std < lower]
+        ratio[input_std > upper] = torch.sigmoid((input_std - upper))[input_std > upper]
+        ratio = ratio  #.detach()
+
+        return ratio
+
+    def forward(self, x):
+
+        # with torch.no_grad():
+        feats_x = self.forward_once(x)
+
+        # print(feats_x.shape)
+        # feats_x = self.generate_alpha(feats_x*255)
+        # pad = nn.ReflectionPad2d(0)
+        # ps_x = self.compute_prob(feats_x)
+        # # print(ps_x[1].shape)
+        # weight = []
+        # for k in range(0, len(self.chns)):
+        #     weight.append(self.entropy(feats_x[k]))
+        # weight = torch.concat(weight, dim=1)
+        # # print(weight.shape)
+        #
+        # weight = weight / weight.sum(dim=(1, 2), keepdim=True)
+        # weight_mean = weight.mean(dim=(1, 2), keepdim=True)
+        # weight_std = torch.sqrt(((weight - weight_mean) ** 2).mean(dim=(1, 2), keepdim=True))
+        # weight = weight.clamp(min=weight_mean - 0.5 * weight_std, max=weight_mean + 0.5 * weight_std)
+        # weight = weight / weight.sum(dim=(1, 2), keepdim=True)
+        # weight_list = torch.split(weight, self.chns, dim=1)
+        # # print(weight_list[1].shape)
+        #
+        # # for k in range(len(self.chns) - 1, -1, -1):
+        # feat_x = F.normalize(feats_x[1], dim=(2, 3))
+        # try:
+        #     x_mean = F.conv2d(pad(feat_x), self.windows[1], stride=1, padding=0, groups=self.chns[1])
+        #     # y_mean = F.conv2d(pad(feat_y), self.windows[k], stride=1, padding=0, groups=self.chns[k])
+        #     x_var = F.conv2d(pad(feat_x ** 2), self.windows[1], stride=1, padding=0,
+        #                      groups=self.chns[1]) - x_mean ** 2
+        #     # y_var = F.conv2d(pad(feat_y ** 2), self.windows[k], stride=1, padding=0,
+        #     #                  groups=self.chns[k]) - y_mean ** 2
+        #     xy_cov = F.conv2d(pad(feat_x * feat_x), self.windows[1], stride=1, padding=0, groups=self.chns[1]) \
+        #              - x_mean * x_mean
+        # except:
+        #     x_mean = feat_x.mean([2, 3], keepdim=True)
+        #     # y_mean = feat_y.mean([2, 3], keepdim=True)
+        #     x_var = ((feat_x - x_mean) ** 2).mean([2, 3], keepdim=True)
+        #     # y_var = ((feat_y - y_mean) ** 2).mean([2, 3], keepdim=True)
+        #     xy_cov = (feat_x * feat_x).mean([2, 3], keepdim=True) - x_mean * x_mean
+        #
+        # T = (2 * x_mean * x_mean + 1e-6) / (x_mean ** 2 + x_mean ** 2 + 1e-6)
+        # S = (2 * xy_cov + 1e-6) / (x_var + x_var + 1e-6)
+        # S = F.interpolate(S, size=(ps_x[1].shape[-2], ps_x[1].shape[-1]), mode='bilinear', align_corners=True)
+        # T = F.interpolate(T, size=(ps_x[1].shape[-2], ps_x[1].shape[-1]), mode='bilinear', align_corners=True)
+        # # print(S.shape, T.shape, ps_x[1].shape)
+        # ps = ps_x[1]  #.expand(x_mean.shape[0], x_mean.shape[1], -1, -1)
+        # pt = 1 - ps
+        # D_map = (pt * T + ps * S) * weight_list[1].unsqueeze(3)
+        # print(D_map.mean([2, 3]))
+        return feats_x
+
+# def prepare_image(image, resize=True):
+#     if resize and min(image.size) > 256:
+#         image = transforms.functional.resize(image, 256)
+#     image = transforms.ToTensor()(image)
+#     return image.unsqueeze(0)
+#
+#
+# if __name__ == '__main__':
+#     from PIL import Image
+#     import argparse
+#
+#     parser = argparse.ArgumentParser()
+#     parser.add_argument('--ref', type=str, default='/data1/shangwei/dataset/video/REDS/val/val_sharp_bicubic/X4/002/00000002.png')
+#     parser.add_argument('--dist', type=str, default='/data1/shangwei/dataset/video/REDS/val/val_sharp_bicubic/X4/002/00000003.png')
+#     args = parser.parse_args()
+#
+#     ref = prepare_image(Image.open(args.ref).convert("RGB"), resize=False)
+#     dist = prepare_image(Image.open(args.dist).convert("RGB"), resize=False)
+#     assert ref.shape == dist.shape
+#     # print(ref.shape)
+#     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+#     model = ADISTS().to(device)
+#     ref = ref.to(device)
+#     dist = dist.to(device)
+#     map = model(ref).mean(1)
+#     # print(score.item())
+#     print(map.shape)
+#     # print(map.min(), map.max())
+#     save = map.cpu().numpy().squeeze() * 255.
+#     save = save.astype(np.uint8)
+#     cv2.imwrite("/data1/shangwei/dataset/video/REDS/val/002_00002_lr_ratio_w21_prob_flat255.png", save)
+

asvr_module/SIREN.py

@@ -0,0 +1,70 @@
+import torch
+from torch import nn
+import numpy as np
+
+class SineLayer(nn.Module):
+    # See paper sec. 3.2, final paragraph, and supplement Sec. 1.5 for discussion of omega_0.
+    
+    # If is_first=True, omega_0 is a frequency factor which simply multiplies the activations before the 
+    # nonlinearity. Different signals may require different omega_0 in the first layer - this is a 
+    # hyperparameter.
+    
+    # If is_first=False, then the weights will be divided by omega_0 so as to keep the magnitude of 
+    # activations constant, but boost gradients to the weight matrix (see supplement Sec. 1.5)
+    
+    def __init__(self, in_features, out_features, bias=True,
+                 is_first=False, omega_0=30):
+        super().__init__()
+        self.omega_0 = omega_0
+        self.is_first = is_first
+        
+        self.in_features = in_features
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        
+        self.init_weights()
+    
+    def init_weights(self):
+        with torch.no_grad():
+            if self.is_first:
+                self.linear.weight.uniform_(-1 / self.in_features, 
+                                             1 / self.in_features)      
+            else:
+                self.linear.weight.uniform_(-np.sqrt(6 / self.in_features) / self.omega_0, 
+                                             np.sqrt(6 / self.in_features) / self.omega_0)
+        
+    def forward(self, input):
+        return torch.sin(self.omega_0 * self.linear(input))
+
+    
+    
+class Siren(nn.Module):
+    def __init__(self, in_features, hidden_features, hidden_layers, out_features, outermost_linear=False, 
+                 first_omega_0=30, hidden_omega_0=30.):
+        super().__init__()
+        
+        self.net = []
+        self.net.append(SineLayer(in_features, hidden_features[0], 
+                                  is_first=True, omega_0=first_omega_0))
+
+        for i in range(hidden_layers):
+            self.net.append(SineLayer(hidden_features[i], hidden_features[i + 1], 
+                                      is_first=False, omega_0=hidden_omega_0))
+
+        if outermost_linear:
+            final_linear = nn.Linear(hidden_features[-1], out_features)
+            
+            with torch.no_grad():
+                final_linear.weight.uniform_(-np.sqrt(6 / hidden_features[-1]) / hidden_omega_0, 
+                                              np.sqrt(6 / hidden_features[-1]) / hidden_omega_0)
+                
+            self.net.append(final_linear)
+        else:
+            self.net.append(SineLayer(hidden_features[-1], out_features, 
+                                      is_first=False, omega_0=hidden_omega_0))
+        
+        self.net = nn.Sequential(*self.net)
+    
+    def forward(self, x):
+        shape = x.shape[:-1]
+        output = self.net(x.view(-1, x.shape[-1]))
+        return output.view(*shape, -1)

asvr_module/asvr_model.py

@@ -0,0 +1,88 @@
+import torch
+# import tqdm
+import torch.nn.functional as F
+
+from .refsrrnn_adists_fgda_only_future import RefsrRNN
+
+
+class ASVR_Config:
+    def __init__(
+        self, model_path: str, space_scale=(8, 8), border: bool = False
+    ) -> None:
+        self.space_scale = space_scale
+        self.model_path = model_path
+        self.border = border
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class ASVR:
+    def __init__(self, cfg: ASVR_Config) -> None:
+        self.net = RefsrRNN()
+        self.net.load_state_dict(torch.load(cfg.model_path))
+        self.net = self.net.to(cfg.device)
+        self.net.eval()
+        self.scale = cfg.space_scale
+        self.device = cfg.device
+
+    def inference(self, frames: torch.Tensor):
+        # frames: [B, H, W, C]
+        print(f"{self.scale=}")
+        print(f"{frames.shape=}")
+        frames = frames.to(self.device)
+        with torch.no_grad():
+            scale_h = torch.ones(1).to(self.device) * (1 / float(self.scale[0]))
+            scale_w = torch.ones(1).to(self.device) * (1 / float(self.scale[1]))
+
+            hs, hw = 1.0 / scale_h, 1.0 / scale_w
+            hs, hw = hs.unsqueeze(-1), hw.unsqueeze(-1)
+            kernel = None
+
+            b, h, w, c = frames.shape
+
+            hr_coord = make_coord((h, w)).unsqueeze(0).to(self.device)
+
+            cell = torch.ones(2).unsqueeze(0).to(self.device)
+            cell[:, 0] *= 2.0 / h
+            cell[:, 1] *= 2.0 / w
+            # [B, C, H, W]
+            frames = frames.permute(0, 3, 1, 2)
+
+            frames_scaled: list[torch.Tensor] = []
+            for frame in frames:
+                frames_scaled.append(
+                    F.interpolate(
+                        frame.unsqueeze(0),
+                        (int(scale_h[0].item() * h), int(scale_w[0].item() * w)),
+                        mode="bicubic",
+                    )
+                )
+
+            if kernel is None:
+                res = torch.zeros(
+                    (
+                        1,
+                        self.net.num_channels,
+                        frames_scaled[0].shape[-2],
+                        frames_scaled[0].shape[-1],
+                    ),
+                    device=frames_scaled[0].device,
+                )
+                kernel = self.net.kernel_predict(res, hr_coord, cell)
+
+            output = self.net.test_forward(frames_scaled, kernel, hr_coord).squeeze(0) #T,C,H,W
+            output = output.permute(0, 2, 3, 1)
+        return output
+
+
+def make_coord(shape):
+    """Make coordinates at grid centers."""
+    coord_seqs = []
+    for i, n in enumerate(shape):
+        # v0, v1 = -1, 1
+
+        r = 1 / n
+        seq = -1 + r + (2 * r) * torch.arange(n).float()
+        coord_seqs.append(seq)
+    # ret = torch.stack(torch.meshgrid(coord_seqs, indexing='ij'), dim=-1)
+    ret = torch.stack(torch.meshgrid(coord_seqs, indexing='ij'), dim=-1)
+    return ret

asvr_module/blocks.py

@@ -0,0 +1,758 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+###############################
+# common
+###############################
+class _NonLocalBlockND(nn.Module):
+    def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True):
+        super(_NonLocalBlockND, self).__init__()
+
+        assert dimension in [1, 2, 3]
+
+        self.dimension = dimension
+        self.sub_sample = sub_sample
+
+        self.in_channels = in_channels
+        self.inter_channels = inter_channels
+
+        if self.inter_channels is None:
+            self.inter_channels = in_channels // 2
+            if self.inter_channels == 0:
+                self.inter_channels = 1
+
+        if dimension == 3:
+            conv_nd = nn.Conv3d
+            max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))
+            bn = nn.BatchNorm3d
+        elif dimension == 2:
+            conv_nd = nn.Conv2d
+            max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))
+            bn = nn.BatchNorm2d
+        else:
+            conv_nd = nn.Conv1d
+            max_pool_layer = nn.MaxPool1d(kernel_size=(2))
+            bn = nn.BatchNorm1d
+
+        self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
+                         kernel_size=1, stride=1, padding=0)
+
+        if bn_layer:
+            self.W = nn.Sequential(
+                conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
+                        kernel_size=1, stride=1, padding=0),
+                bn(self.in_channels)
+            )
+            nn.init.constant_(self.W[1].weight, 0)
+            nn.init.constant_(self.W[1].bias, 0)
+        else:
+            self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
+                             kernel_size=1, stride=1, padding=0)
+            nn.init.constant_(self.W.weight, 0)
+            nn.init.constant_(self.W.bias, 0)
+
+        self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
+                             kernel_size=1, stride=1, padding=0)
+
+        self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
+                           kernel_size=1, stride=1, padding=0)
+
+        # self.fusion = conv_nd(in_channels=self.inter_channels, out_channels=self.inter_channels,
+        #                    kernel_size=1, stride=1, padding=0)
+
+        if sub_sample:
+            self.g = nn.Sequential(self.g, max_pool_layer)
+            self.phi = nn.Sequential(self.phi, max_pool_layer)
+
+        self.softmax = nn.Softmax(dim=2)
+
+    def forward(self, x, eve):
+        '''
+        :param x: (b, c, t, h, w)
+        :return:
+        '''
+        batch_size = x.size(0)
+
+        g_x = self.g(x).view(batch_size, self.inter_channels, -1)   #b, c, h*w/4
+        g_x = g_x.permute(0, 2, 1)
+
+        theta_eve = self.theta(eve).view(batch_size, self.inter_channels, -1)   #b, c, h*w
+        theta_eve = theta_eve.permute(0, 2, 1)
+        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)   #b, c, h*w/4
+        f = torch.matmul(theta_eve, phi_x)   # (b, h*w, c)  X  (b, c, h*w/4)   =  (b, h*w, h*w/4)
+        N = f.size(-1)
+        f_div_C = f / N
+        f_div_C = self.softmax(f_div_C)
+
+        y = torch.matmul(f_div_C, g_x)    # (b, h*w, h*w/4) X (b, h*w/4, c) = (b, h*w, c)
+        y = y.permute(0, 2, 1).contiguous()
+        y = y.view(batch_size, self.inter_channels, *x.size()[2:])   # b,c,h,w
+        W_y = self.W(y)
+        z = W_y + x
+
+        return z
+
+
+class NONLocalBlock2D(_NonLocalBlockND):
+    def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=False):
+        super(NONLocalBlock2D, self).__init__(in_channels,
+                                              inter_channels=inter_channels,
+                                              dimension=2, sub_sample=sub_sample,
+                                              bn_layer=bn_layer)
+
+
+def get_same_padding(kernel_size, dilation):
+    kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1)
+    padding = (kernel_size - 1) // 2
+    return padding
+
+class ConvLSTM(nn.Module):
+    def __init__(self, inp_dim, oup_dim, kernel, dilation):
+        super().__init__()
+        pad_x = int(dilation * (kernel - 1) / 2)
+        self.conv_xf = nn.Conv2d(inp_dim, oup_dim, kernel, padding=pad_x, dilation=dilation)
+        self.conv_xi = nn.Conv2d(inp_dim, oup_dim, kernel, padding=pad_x, dilation=dilation)
+        self.conv_xo = nn.Conv2d(inp_dim, oup_dim, kernel, padding=pad_x, dilation=dilation)
+        self.conv_xj = nn.Conv2d(inp_dim, oup_dim, kernel, padding=pad_x, dilation=dilation)
+
+        pad_h = int((kernel - 1) / 2)
+        self.conv_hf = nn.Conv2d(oup_dim, oup_dim, kernel, padding=pad_h)
+        self.conv_hi = nn.Conv2d(oup_dim, oup_dim, kernel, padding=pad_h)
+        self.conv_ho = nn.Conv2d(oup_dim, oup_dim, kernel, padding=pad_h)
+        self.conv_hj = nn.Conv2d(oup_dim, oup_dim, kernel, padding=pad_h)
+
+        self.se = SEBlock(oup_dim, 4)
+        self.relu = nn.LeakyReLU(0.2)
+
+    def forward(self, x, pair=None):
+        if pair is None:
+            i = F.sigmoid(self.conv_xi(x))
+            o = F.sigmoid(self.conv_xo(x))
+            j = F.tanh(self.conv_xj(x))
+            c = i * j
+            h = o * c
+        else:
+            h, c = pair
+            f = F.sigmoid(self.conv_xf(x) + self.conv_hf(h))
+            i = F.sigmoid(self.conv_xi(x) + self.conv_hi(h))
+            o = F.sigmoid(self.conv_xo(x) + self.conv_ho(h))
+            j = F.tanh(self.conv_xj(x) + self.conv_hj(h))
+            c = f * c + i * j
+            h = o * F.tanh(c)
+
+        h = self.relu(self.se(h))
+        return h, [h, c]
+
+class SEBlock(nn.Module):
+    def __init__(self, input_dim, reduction):
+        super().__init__()
+        mid = int(input_dim / reduction)
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(input_dim, mid),
+            nn.ReLU(inplace=True),
+            nn.Linear(mid, input_dim),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y
+
+###############################
+# ResNet
+###############################
+# class ResBlock(nn.Module):
+#     def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bn=False, se=True):
+#         super(ResBlock, self).__init__()
+#         self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, stride=stride,
+#                                padding=get_same_padding(kernel_size, dilation), dilation=dilation)
+#         self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, stride=1,
+#                                padding=get_same_padding(kernel_size, dilation), dilation=dilation)
+#         self.relu = nn.LeakyReLU(0.1, inplace=True)
+#         if se:
+#             self.se = SEBlock(planes, 4)
+#         if bn:
+#             self.bn1 = nn.InstanceNorm2d(planes)
+#             self.bn2 = nn.InstanceNorm2d(planes)
+#         self.se_ = se
+#         self.bn_ = bn
+#         self.res_translate = None
+#         if not inplanes == planes or not stride == 1:
+#             self.res_translate = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride)
+#
+#     def forward(self, x):
+#         residual = x
+#         if self.bn_:
+#             out = self.relu(self.bn1(self.conv1(x)))
+#         else:
+#             out = self.relu(self.conv1(x))
+#         if self.bn_:
+#             out = self.bn2(self.conv2(out))
+#         else:
+#             out = self.conv2(out)
+#         if self.se_:
+#             out = self.se(out)
+#         if self.res_translate is not None:
+#             residual = self.res_translate(residual)
+#         out += residual
+#
+#         return out
+
+class CoralLayer(torch.nn.Module):
+    """ Implements CORAL layer described in
+
+    Cao, Mirjalili, and Raschka (2020)
+    *Rank Consistent Ordinal Regression for Neural Networks
+       with Application to Age Estimation*
+    Pattern Recognition Letters, https://doi.org/10.1016/j.patrec.2020.11.008
+
+    Parameters
+    -----------
+    size_in : int
+        Number of input features for the inputs to the forward method, which
+        are expected to have shape=(num_examples, num_features).
+
+    num_classes : int
+        Number of classes in the dataset.
+
+
+    """
+    def __init__(self, size_in, num_classes):
+        super().__init__()
+        self.size_in, self.size_out = size_in, 1
+
+        self.coral_weights = torch.nn.Linear(self.size_in, 1, bias=False)
+        self.coral_bias = torch.nn.Parameter(
+             torch.zeros(num_classes-1).float())
+
+    def forward(self, x):
+        """
+        Computes forward pass.
+
+        Parameters
+        -----------
+        x : torch.tensor, shape=(num_examples, num_features)
+            Input features.
+
+        Returns
+        -----------
+        logits : torch.tensor, shape=(num_examples, num_classes-1)
+        """
+        return self.coral_weights(x) + self.coral_bias
+
+class ResBlock(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bn=False, se=True):
+        super(ResBlock, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, stride=stride,
+                               padding=get_same_padding(kernel_size, dilation), dilation=dilation)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, stride=1,
+                               padding=get_same_padding(kernel_size, dilation), dilation=dilation)
+        self.relu = nn.LeakyReLU(0.1, inplace=True)
+        if se:
+            self.se = SEBlock(planes, 4)
+        if bn:
+            self.bn1 = nn.BatchNorm2d(planes)
+            self.bn2 = nn.BatchNorm2d(planes)
+        self.se_ = se
+        self.bn_ = bn
+        self.res_translate = None
+        if not inplanes == planes or not stride == 1:
+            self.res_translate = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride)
+
+    def forward(self, x):
+        residual = x
+        if self.bn_:
+            out = self.relu(self.bn1(self.conv1(x)))
+        else:
+            out = self.relu(self.conv1(x))
+        if self.bn_:
+            out = self.bn2(self.conv2(out))
+        else:
+            out = self.conv2(out)
+        if self.se_:
+            out = self.se(out)
+        if self.res_translate is not None:
+            residual = self.res_translate(residual)
+        out += residual
+
+        return out
+
+class GResBlock(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1,group=1):
+        super(GResBlock, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, stride=stride,
+                               padding=get_same_padding(kernel_size, dilation), dilation=dilation, groups=group)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, stride=1,
+                               padding=get_same_padding(kernel_size, dilation), dilation=dilation, groups=group)
+        self.relu = nn.LeakyReLU(0.1, inplace=True)
+        self.conv3 = nn.Conv2d(planes, planes, kernel_size=kernel_size, stride=1,
+                               padding=get_same_padding(kernel_size, dilation), dilation=dilation, groups=group)
+        self.res_translate = None
+        if not inplanes == planes or not stride == 1:
+            self.res_translate = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride)
+
+    def forward(self, x):
+        residual = x
+
+        out = self.relu(self.conv1(x))
+        out = self.relu(self.conv2(out))
+        out = self.conv3(out)
+        if self.res_translate is not None:
+            residual = self.res_translate(residual)
+        out += residual
+
+        return out
+
+class SKConv(nn.Module):
+    def __init__(self, features, WH, M, G, r, stride=1, L=32):
+        """ Constructor
+        Args:
+            features: input channel dimensionality.
+            WH: input spatial dimensionality, used for GAP kernel size.
+            M: the number of branchs.
+            G: num of convolution groups.
+            r: the radio for compute d, the length of z.
+            stride: stride, default 1.
+            L: the minimum dim of the vector z in paper, default 32.
+        """
+        super(SKConv, self).__init__()
+        d = max(int(features / r), L)
+        self.M = M
+        self.features = features
+        self.convs = nn.ModuleList([])
+        for i in range(M):
+            self.convs.append(nn.Sequential(
+                nn.Conv2d(features, features, kernel_size=3 + i * 2, stride=stride, padding=1 + i, groups=G),
+                # nn.BatchNorm2d(features),
+                nn.LeakyReLU(0.1)
+            ))
+        # self.gap = nn.AvgPool2d(int(WH/stride))
+        self.fc = nn.Linear(features, d)
+        self.fcs = nn.ModuleList([])
+        for i in range(M):
+            self.fcs.append(
+                nn.Linear(d, features)
+            )
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+        for i, conv in enumerate(self.convs):
+            fea = conv(x).unsqueeze_(dim=1)
+            if i == 0:
+                feas = fea
+            else:
+                feas = torch.cat([feas, fea], dim=1)
+        fea_U = torch.sum(feas, dim=1)
+        # fea_s = self.gap(fea_U).squeeze_()
+        fea_s = fea_U.mean(-1).mean(-1)
+        fea_z = self.fc(fea_s)
+        for i, fc in enumerate(self.fcs):
+            vector = fc(fea_z).unsqueeze_(dim=1)
+            if i == 0:
+                attention_vectors = vector
+            else:
+                attention_vectors = torch.cat([attention_vectors, vector], dim=1)
+        attention_vectors = self.softmax(attention_vectors)
+        attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1)
+        fea_v = (feas * attention_vectors).sum(dim=1)
+        return fea_v
+
+class BasicConv(nn.Module):
+    def __init__(self, in_channel, out_channel, kernel_size, stride, bias=True, norm=False, relu=True, transpose=False):
+        super(BasicConv, self).__init__()
+        if bias and norm:
+            bias = False
+
+        padding = kernel_size // 2
+        layers = list()
+        if transpose:
+            padding = kernel_size // 2 -1
+            layers.append(nn.ConvTranspose2d(in_channel, out_channel, kernel_size, padding=padding, stride=stride, bias=bias))
+        else:
+            layers.append(
+                nn.Conv2d(in_channel, out_channel, kernel_size, padding=padding, stride=stride, bias=bias))
+        if norm:
+            layers.append(nn.BatchNorm2d(out_channel))
+        if relu:
+            layers.append(nn.ReLU(inplace=True))
+        self.main = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.main(x)
+
+class SCM(nn.Module):
+    def __init__(self, out_plane):
+        super(SCM, self).__init__()
+        self.main = nn.Sequential(
+            BasicConv(3, out_plane//4, kernel_size=3, stride=1, relu=True),
+            BasicConv(out_plane // 4, out_plane // 2, kernel_size=1, stride=1, relu=True),
+            BasicConv(out_plane // 2, out_plane // 2, kernel_size=3, stride=1, relu=True),
+            BasicConv(out_plane // 2, out_plane-3, kernel_size=1, stride=1, relu=True)
+        )
+
+        self.conv = BasicConv(out_plane, out_plane, kernel_size=1, stride=1, relu=False)
+
+    def forward(self, x):
+        x = torch.cat([x, self.main(x)], dim=1)
+        return self.conv(x)
+
+class SKUnit(nn.Module):
+    def __init__(self, in_features, out_features, WH, M, G, r, mid_features=None, stride=1, L=32):
+        """ Constructor
+        Args:
+            in_features: input channel dimensionality.
+            out_features: output channel dimensionality.
+            WH: input spatial dimensionality, used for GAP kernel size.
+            M: the number of branchs.
+            G: num of convolution groups.
+            r: the radio for compute d, the length of z.
+            mid_features: the channle dim of the middle conv with stride not 1, default out_features/2.
+            stride: stride.
+            L: the minimum dim of the vector z in paper.
+        """
+        super(SKUnit, self).__init__()
+        if mid_features is None:
+            mid_features = int(out_features / 2)
+        self.feas = nn.Sequential(
+            nn.Conv2d(in_features, mid_features, 1, stride=1),
+            # nn.BatchNorm2d(mid_features),
+            SKConv(mid_features, WH, M, G, r, stride=stride, L=L),
+            # nn.BatchNorm2d(mid_features),
+            nn.Conv2d(mid_features, out_features, 1, stride=1),
+            # nn.BatchNorm2d(out_features)
+        )
+        if in_features == out_features:  # when dim not change, in could be added diectly to out
+            self.shortcut = nn.Sequential()
+        else:  # when dim not change, in should also change dim to be added to out
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_features, out_features, 1, stride=stride),
+                # nn.BatchNorm2d(out_features)
+            )
+
+    def forward(self, x):
+        fea = self.feas(x)
+        return fea + self.shortcut(x)
+
+
+class ChannelAttention(nn.Module):
+    def __init__(self, in_planes, ratio=16):
+        super(ChannelAttention, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc = nn.Sequential(nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
+                                nn.ReLU(),
+                                nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False))
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.fc(self.avg_pool(x))
+        max_out = self.fc(self.max_pool(x))
+        out = avg_out + max_out
+        return self.sigmoid(out)
+
+
+class SpatialAttention(nn.Module):
+    def __init__(self, kernel_size=7):
+        super(SpatialAttention, self).__init__()
+
+        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        x = self.conv1(x)
+        return self.sigmoid(x)
+
+# class Unet(nn.Module):
+#
+#     def __init__(self, in_dim=3*4, out_dim=3*3, block=2, feats=32, kernel_size=3):
+#         super(Unet, self).__init__()
+#         InBlock = []
+#
+#         InBlock.extend([nn.Sequential(
+#             nn.Conv2d(in_dim, feats, kernel_size=7, stride=1,
+#                       padding=7 // 2),
+#             nn.LeakyReLU(0.1,inplace=True)
+#         )])
+#
+#         InBlock.extend([ResBlock(feats, feats, kernel_size=kernel_size, stride=1)
+#                         for _ in range(block)])
+#
+#         # encoder1
+#         Encoder_first = [nn.Sequential(
+#             nn.Conv2d(feats, feats*2, kernel_size=kernel_size, stride=2, padding=kernel_size // 2),
+#             nn.LeakyReLU(inplace=True)
+#         )]
+#         Encoder_first.extend([ResBlock(feats*2, feats*2, kernel_size=kernel_size, stride=1)
+#                               for _ in range(block)])
+#         # encoder2
+#         Encoder_second = [nn.Sequential(
+#             nn.Conv2d(feats*2, feats*4, kernel_size=kernel_size, stride=2, padding=kernel_size // 2),
+#             nn.LeakyReLU(inplace=True)
+#         )]
+#         Encoder_second.extend([ResBlock(feats * 4, feats * 4, kernel_size=kernel_size, stride=1)
+#                                for _ in range(block)])
+#
+#         # decoder2
+#         Decoder_second = [ResBlock(feats * 4, feats * 4, kernel_size=kernel_size, stride=1)
+#                           for _ in range(block)]
+#         Decoder_second.append(nn.Sequential(
+#             nn.ConvTranspose2d(feats * 4, feats * 2, kernel_size=3, stride=2, padding=1, output_padding=1),
+#             nn.LeakyReLU(inplace=True)
+#         ))
+#         # decoder1
+#         Decoder_first = [ResBlock(feats * 2, feats * 2, kernel_size=kernel_size, stride=1)
+#                          for _ in range(block)]
+#         Decoder_first.append(nn.Sequential(
+#             nn.ConvTranspose2d(feats * 2, feats, kernel_size=3, stride=2, padding=1, output_padding=1),
+#             nn.LeakyReLU(inplace=True)
+#         ))
+#
+#         OutBlock = [ResBlock(feats, feats, kernel_size=kernel_size, stride=1)
+#                     for _ in range(block)]
+#         OutBlock.append(nn.Conv2d(feats, out_dim, kernel_size=kernel_size, stride=1, padding=kernel_size // 2))
+#
+#
+#         self.inBlock = nn.Sequential(*InBlock)
+#         self.encoder_first = nn.Sequential(*Encoder_first)
+#         self.encoder_second = nn.Sequential(*Encoder_second)
+#         self.decoder_second = nn.Sequential(*Decoder_second)
+#         self.decoder_first = nn.Sequential(*Decoder_first)
+#         self.outBlock = nn.Sequential(*OutBlock)
+#
+#     def forward(self, x):
+#
+#         x = torch.cat(x, dim=1)
+#
+#         first_scale_inblock = self.inBlock(x)
+#         first_scale_encoder_first = self.encoder_first(first_scale_inblock)
+#         first_scale_encoder_second = self.encoder_second(first_scale_encoder_first)
+#
+#         return first_scale_inblock, first_scale_encoder_first, first_scale_encoder_second
+
+
+class DA_conv(nn.Module):
+    def __init__(self, channels_in, channels_out, kernel_size, reduction, kernel_dim):
+        super(DA_conv, self).__init__()
+        self.channels_out = channels_out
+        self.channels_in = channels_in
+        self.kernel_size = kernel_size
+
+        # self.kernel = nn.Sequential(
+        #     nn.Linear(kernel_dim, kernel_dim//2, bias=False),
+        #     nn.LeakyReLU(0.1, True),
+        #     nn.Linear(kernel_dim//2, channels_in * self.kernel_size * self.kernel_size, bias=False)  #3
+        # )
+        # self.conv = nn.Conv2d(channels_in, channels_out, 1, padding=0)
+        self.ca = CA_layer(kernel_dim, channels_out, reduction)
+
+        # self.relu = nn.LeakyReLU(0.1, True)
+
+    def forward(self, x, vec):
+        '''
+        :param x[0]: feature map: B * C * H * W
+        :param x[1]: degradation representation: B * C  #4,256
+        '''
+        # print(x.size(), vec.size())
+        # b, c, h, w = x.size()
+        # # print(x[0].size())
+        # # branch 1
+        # kernel = self.kernel(vec).view(-1, 1, self.kernel_size, self.kernel_size)  #4,128*3*3   512,1,3,3
+        # # print(kernel.size())
+        # out = self.relu(F.conv2d(x.view(1, -1, h, w), kernel, groups=b*c, padding=(self.kernel_size-1)//2))
+        # out = self.conv(out.view(b, -1, h, w))
+        # print(out.size())
+        # branch 2
+        out = x
+        out = out + self.ca(x, vec)
+
+        return out
+
+
+class CA_layer(nn.Module):
+    def __init__(self, channels_in, channels_out, reduction):
+        super(CA_layer, self).__init__()
+        self.conv_du = nn.Sequential(
+            nn.Conv2d(channels_in, channels_in//reduction, 1, 1, 0, bias=False),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(channels_in // reduction, channels_out, 1, 1, 0, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x,vec):
+        '''
+        :param x[0]: feature map: B * C * H * W
+        :param x[1]: degradation representation: B * C
+        '''
+        att = self.conv_du(vec[:, :, None, None])
+
+        return x * att
+
+class ChanNorm(nn.Module):
+    def __init__(self, dim, eps = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
+        self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
+
+    def forward(self, x):
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
+        mean = torch.mean(x, dim = 1, keepdim = True)
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
+# nn.InstanceNorm2d
+# b=torch.tensor([
+#     [
+#         [[1,0,1],[0,1,0],[1,0,1]],[[1,0,1],[0,1,0],[1,0,1]]
+#     ],
+#     [
+#         [[1, 0, 1], [0, 1, 0], [1, 0, 1]], [[1, 0, 1], [0, 1, 0], [1, 0, 1]]
+#     ]
+#     ])
+
+class Conv2DMod(nn.Module):
+    def __init__(self, in_chan, out_chan, kernel, demod=True, stride=1, dilation=1, eps = 1e-8, **kwargs):
+        super().__init__()
+        self.filters = out_chan
+        self.demod = demod
+        self.kernel = kernel
+        self.stride = stride
+        self.dilation = dilation
+        self.weight = nn.Parameter(torch.randn((out_chan, in_chan, kernel, kernel)))
+        self.eps = eps
+        nn.init.kaiming_normal_(self.weight, a=0, mode='fan_in', nonlinearity='leaky_relu')
+
+    def _get_same_padding(self, size, kernel, dilation, stride):
+        return ((size - 1) * (stride - 1) + dilation * (kernel - 1)) // 2
+
+    def forward(self, x, y):
+        b, c, h, w = x.shape
+
+        w1 = y[:, None, :, None, None]
+        w2 = self.weight[None, :, :, :, :]
+        weights = w2 * (w1 + 1)
+
+        if self.demod:
+            d = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + self.eps)
+            weights = weights * d
+
+        x = x.reshape(1, -1, h, w)
+
+        _, _, *ws = weights.shape
+        weights = weights.reshape(b * self.filters, *ws)
+
+        padding = self._get_same_padding(h, self.kernel, self.dilation, self.stride)
+        x = F.conv2d(x, weights, padding=padding, groups=b)
+
+        x = x.reshape(-1, self.filters, h, w)
+        return x
+
+class Sty_layer(nn.Module):
+    def __init__(self, channels_in, channels_out, reduction):
+        super(Sty_layer, self).__init__()
+        self.norm = ChanNorm(channels_out)
+        self.conv_du = nn.Sequential(
+            nn.Linear(channels_in, channels_in//reduction, bias=False),
+            nn.LeakyReLU(0.1, True),
+            nn.Linear(channels_in // reduction, channels_out, bias=False),
+            # nn.Sigmoid()
+        )
+        self.conv = Conv2DMod(channels_out, channels_out, 3)
+
+    def forward(self, x, vec):
+        '''
+        :param x[0]: feature map: B * C * H * W
+        :param x[1]: degradation representation: B * C
+        '''
+        norm_x = self.norm(x)
+        sty = self.conv_du(vec)
+        out = self.conv(norm_x, sty)
+
+        return out + x
+
+
+class Ada_layer(nn.Module):
+    def __init__(self, channels_in, channels_out, reduction):
+        super(Ada_layer, self).__init__()
+        self.conv_alpha = nn.Sequential(
+            nn.Conv2d(channels_in, channels_in//reduction, 1, 1, 0, bias=False),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(channels_in // reduction, channels_out, 1, 1, 0, bias=False),
+            nn.Sigmoid()
+        )
+        self.conv_belta = nn.Sequential(
+            nn.Conv2d(channels_in, channels_in // reduction, 1, 1, 0, bias=False),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(channels_in // reduction, channels_out, 1, 1, 0, bias=False)
+        )
+
+    def forward(self, x,vec):
+        '''
+        :param x[0]: feature map: B * C * H * W
+        :param x[1]: degradation representation: B * C
+        '''
+        alpha = self.conv_alpha(vec[:, :, None, None])
+        belta = self.conv_belta(vec[:, :, None, None])
+
+        return x * alpha + belta
+
+class DAB(nn.Module):
+    def __init__(self, n_feat, kernel_size, kernel_dim, reduction=8):
+        super(DAB, self).__init__()
+
+        self.da_conv1 = DA_conv(n_feat, n_feat, kernel_size, reduction, kernel_dim)
+        self.da_conv2 = DA_conv(n_feat, n_feat, kernel_size, reduction, kernel_dim)
+        self.conv1 = nn.Conv2d(n_feat, n_feat, kernel_size, padding=(kernel_size//2))
+        self.conv2 = nn.Conv2d(n_feat, n_feat, kernel_size, padding=(kernel_size//2))
+
+        self.relu =  nn.LeakyReLU(0.1, True)
+
+    def forward(self, inp):
+        '''
+        :param x[0]: feature map: B * C * H * W
+        :param x[1]: degradation representation: B * C
+        '''
+        x, vec = inp
+        out = self.relu(self.da_conv1(x,vec))
+        out = self.relu(self.conv1(out))
+        out = self.relu(self.da_conv2(out, vec))
+        out = self.conv2(out) + x
+
+        return out,vec
+
+class SFT_Layer(nn.Module):
+    def __init__(self, nf=64, para=10):
+        super(SFT_Layer, self).__init__()
+        self.mul_conv1 = nn.Conv2d(para + nf, nf//2, kernel_size=3, stride=1, padding=1)
+        self.mul_leaky = nn.LeakyReLU(0.2)
+        self.mul_conv2 = nn.Conv2d(nf//2, nf, kernel_size=3, stride=1, padding=1)
+
+        self.add_conv1 = nn.Conv2d(para + nf, nf//2, kernel_size=3, stride=1, padding=1)
+        self.add_leaky = nn.LeakyReLU(0.2)
+        self.add_conv2 = nn.Conv2d(nf//2, nf, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, feature_maps, para_maps):
+        cat_input = torch.cat((feature_maps, para_maps), dim=1)
+        mul = torch.sigmoid(self.mul_conv2(self.mul_leaky(self.mul_conv1(cat_input))))
+        add = self.add_conv2(self.add_leaky(self.add_conv1(cat_input)))
+        return feature_maps * mul + add
+
+
+class SFT_Residual_Block(nn.Module):
+    def __init__(self, nf=64, para=10):
+        super(SFT_Residual_Block, self).__init__()
+        self.sft1 = SFT_Layer(nf=nf, para=para)
+        self.sft2 = SFT_Layer(nf=nf, para=para)
+        self.conv1 = nn.Conv2d(in_channels=nf, out_channels=nf, kernel_size=3, stride=1, padding=1, bias=True)
+        self.conv2 = nn.Conv2d(in_channels=nf, out_channels=nf, kernel_size=3, stride=1, padding=1, bias=True)
+
+    def forward(self, feature_maps, para_maps):
+        fea1 = F.relu(self.sft1(feature_maps, para_maps))
+        fea2 = F.relu(self.sft2(self.conv1(fea1), para_maps))
+        fea3 = self.conv2(fea2)
+        return torch.add(feature_maps, fea3)

asvr_module/components.py

@@ -0,0 +1,120 @@
+import functools
+from .init import init_fn
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ResBlock(nn.Module):
+    def __init__(self, in_channels, mid_channels, out_channels):
+        super(ResBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=mid_channels, kernel_size=3, stride=1, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        output = self.conv2(self.relu(self.conv1(x)))
+        output = torch.add(output, x)
+        return output
+
+class ResBlocks(nn.Module):
+    def __init__(self, input_channels, num_resblocks, num_channels):
+        super(ResBlocks, self).__init__()
+        self.input_channels = input_channels
+        self.first_conv = nn.Conv2d(in_channels=self.input_channels, out_channels=num_channels, kernel_size=3, stride=1, padding=1)
+
+        modules = []
+        for _ in range(num_resblocks):
+            modules.append(ResBlock(in_channels=num_channels, mid_channels=num_channels, out_channels=num_channels))
+        self.resblocks = nn.Sequential(*modules)
+
+        fn = functools.partial(init_fn, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2)
+        self.apply(fn)
+
+    def forward(self, h):
+        shallow_feature = self.first_conv(h)
+        new_h = self.resblocks(shallow_feature)
+        return new_h
+
+class Conv2DMod(nn.Module):
+    def __init__(self, in_chan, out_chan, kernel, style_ch=82, demod=True, stride=1, padding=1, dilation=1, eps = 1e-8, **kwargs):
+        super().__init__()
+        self.filters = out_chan
+        self.demod = demod
+        self.kernel = kernel
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.weight = nn.Parameter(torch.randn((out_chan, in_chan, kernel, kernel)))
+        self.eps = eps
+        self.to_style = nn.Linear(style_ch, in_chan)
+        nn.init.kaiming_normal_(self.weight, a=0, mode='fan_in', nonlinearity='leaky_relu')
+
+    # def _get_same_padding(self, size, kernel, dilation, stride):
+    #     return ((size - 1) * (stride - 1) + dilation * (kernel - 1)) // 2
+
+    def forward(self, x, scale):
+        b, c, h, w = x.shape
+
+        y = self.to_style(scale)
+        w1 = y[:, None, :, None, None]
+        # print('w1--', w1.shape)
+        w2 = self.weight[None, :, :, :, :]
+        # print('w2--',w2.shape)
+        weights = w2 * (w1 + 1)
+
+        if self.demod:
+            d = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + self.eps)
+            weights = weights * d
+
+        x = x.reshape(1, -1, h, w)
+
+        _, _, *ws = weights.shape
+        weights = weights.reshape(b * self.filters, *ws)
+
+        # padding = self._get_same_padding(h, self.kernel, self.dilation, self.stride)
+        x = F.conv2d(x, weights, padding=self.padding, groups=b)
+        # print(x.shape)
+
+        x = x.reshape(-1, self.filters, h, w)
+        return x
+
+class ResBlocks_mod(nn.Module):
+    def __init__(self, input_channels, num_resblocks, num_channels):
+        super(ResBlocks_mod, self).__init__()
+        self.input_channels = input_channels
+        self.first_conv = nn.Conv2d(in_channels=self.input_channels, out_channels=num_channels, kernel_size=3, stride=1, padding=1)
+        self.num_resblocks = num_resblocks
+        modules = []
+        for _ in range(num_resblocks):
+            modules.append(ResBlock(in_channels=num_channels, mid_channels=num_channels, out_channels=num_channels))
+        self.resblocks = nn.Sequential(*modules)
+        fn = functools.partial(init_fn, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2)
+        self.apply(fn)
+
+        modc = []
+        for _ in range(num_resblocks // 5):
+            modc.append(Conv2DMod(num_channels, num_channels, 3))
+        self.mbconv = nn.Sequential(*modc)
+
+    def forward(self, h, scale):
+        new_h = self.first_conv(h)
+        for i in range(self.num_resblocks//5):
+            new_h = self.mbconv[i](new_h, scale)
+            new_h = self.resblocks[i*5:(i+1)*5](new_h)
+        return new_h
+
+class D(nn.Module):
+    def __init__(self, in_channels, mid_channels, out_channels):
+        super(D, self).__init__()
+        layers = []
+        layers.append(nn.Conv2d(in_channels=in_channels, out_channels=mid_channels, kernel_size=3, stride=1, padding=1))
+        layers.append(nn.ReLU())
+        layers.append(nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1))
+        self.convs = nn.Sequential(*layers)
+
+        fn = functools.partial(init_fn, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2)
+        self.apply(fn)
+
+    def forward(self, x):
+        x = self.convs(x)
+        return x

asvr_module/cuf_siren.py

@@ -0,0 +1,253 @@
+import numpy as np
+
+import torch
+import torch.fft
+import torch.nn as nn
+import torch.nn.functional as F
+from .SIREN import Siren
+
+def make_coord(shape, ranges=None, flatten=True):
+    """ Make coordinates at grid centers.
+        coord_x = -1+(2*i+1)/W
+        coord_y = -1+(2*i+1)/H
+        normalize to (-1, 1)
+    """
+    coord_seqs = []
+    for i, n in enumerate(shape):
+        if ranges is None:
+            v0, v1 = -1, 1
+        else:
+            v0, v1 = ranges[i]
+        r = (v1 - v0) / (2 * n)
+        seq = v0 + r + (2 * r) * torch.arange(n).float()
+        coord_seqs.append(seq)
+    ret = torch.stack(torch.meshgrid(*coord_seqs, indexing='ij'), dim=-1)
+    if flatten:
+        ret = ret.view(-1, ret.shape[-1])
+    return ret
+
+class PositionEncoder(nn.Module):
+    def __init__(
+        self,
+        posenc_type='sinusoid',
+        complex_transform=False,
+        posenc_scale=10,
+        gauss_scale=1,
+        in_dims=2,
+        enc_dims=64,
+        hidden_dims=64,
+        out_dims=32,
+        head=8,
+        gamma=1
+    ):
+        super().__init__()
+
+        self.posenc_type = posenc_type
+        self.complex_transform = complex_transform
+        self.posenc_scale = posenc_scale
+        self.gauss_scale = gauss_scale
+        self.out_dims = out_dims
+        self.in_dims = in_dims
+        self.enc_dims = enc_dims
+        self.hidden_dims = hidden_dims
+        self.head = head
+        self.gamma = gamma
+
+        self.define_parameter()
+
+    def define_parameter(self):
+        if self.posenc_type == 'sinusoid' or self.posenc_type == 'ipe':
+            self.b_vals = 2.**torch.linspace(
+                0, self.posenc_scale, self.enc_dims // 4
+            ) - 1  # -1 -> (2 * pi)
+            self.b_vals = torch.stack([self.b_vals, torch.zeros_like(self.b_vals)], dim=-1)
+            self.b_vals = torch.cat([self.b_vals, torch.roll(self.b_vals, 1, -1)], dim=0)
+            self.a_vals = torch.ones(self.b_vals.shape[0])
+            self.proj = nn.Linear(self.enc_dims, self.out_dims) #self.head)
+
+        elif self.posenc_type == 'learn':
+            self.Wr = nn.Linear(self.in_dims, self.hidden_dims // 2, bias=False)
+            self.mlp = nn.Sequential(
+                nn.LayerNorm(self.hidden_dims, eps=1e-6),
+                nn.Linear(self.hidden_dims, self.hidden_dims),
+                nn.GELU(),
+                nn.LayerNorm(self.hidden_dims, eps=1e-6),
+                nn.Linear(self.hidden_dims, self.enc_dims)
+            )
+            self.proj = nn.Sequential(nn.GELU(), nn.Linear(self.enc_dims, self.head))
+            self.init_weight()
+
+        elif self.posenc_type == 'dpb':
+            self.mlp = nn.Sequential(
+                nn.Linear(2, self.hidden_dims),
+                nn.LayerNorm(self.hidden_dims, eps=1e-6),
+                nn.ReLU(),
+                nn.Linear(self.hidden_dims, self.hidden_dims),
+                nn.LayerNorm(self.hidden_dims, eps=1e-6),
+                nn.ReLU(),
+                nn.Linear(self.hidden_dims, self.enc_dims)
+            )
+            self.proj = nn.Sequential(
+                nn.LayerNorm(self.hidden_dims, eps=1e-6),
+                nn.ReLU(),
+                nn.Linear(self.enc_dims, self.head)
+            )
+
+    def init_weight(self):
+        nn.init.normal_(self.Wr.weight.data, mean=0, std=self.gamma**-2)
+
+    def forward(self, positions, cells=None):
+
+        if self.posenc_type is None:
+            return positions
+
+        if self.posenc_type == 'sinusoid' or self.posenc_type == 'ipe':
+            self.b_vals = self.b_vals.to(positions.device) #.cuda()
+            self.a_vals = self.a_vals.to(positions.device) #.cuda()
+
+            # b, q, 1, c (x -> c/2, y -> c/2)
+            sin_part = self.a_vals * torch.sin(
+                torch.matmul(positions, self.b_vals.transpose(-2, -1))
+            )
+            cos_part = self.a_vals * torch.cos(
+                torch.matmul(positions, self.b_vals.transpose(-2, -1))
+            )
+
+            if self.posenc_type == 'ipe':
+                # b, q, 2
+                cell = cells.clone()
+                cell_part = torch.sinc(
+                    torch.matmul((1 / np.pi * cell), self.b_vals.transpose(-2, -1))
+                )
+
+                sin_part = sin_part * cell_part
+                cos_part = cos_part * cell_part
+
+            if self.complex_transform:
+                pos_enocoding = torch.view_as_complex(torch.stack([cos_part, sin_part], dim=-1))
+            else:
+                pos_enocoding = torch.cat([sin_part, cos_part], dim=-1)
+                # pos_bias = self.proj(pos_enocoding)
+
+        elif self.posenc_type == 'learn':
+            projected_pos = self.Wr(positions)
+
+            sin_part = torch.sin(projected_pos)
+            cos_part = torch.cos(projected_pos)
+
+            if self.complex_transform:
+                pos_enocoding = 1 / np.sqrt(self.hidden_dims) * torch.view_as_complex(
+                    torch.stack([cos_part, sin_part], dim=-1)
+                )
+            else:
+                pos_enocoding = 1 / np.sqrt(self.hidden_dims
+                                           ) * torch.cat([sin_part, cos_part], dim=-1)
+                pos_enocoding = self.mlp(pos_enocoding)
+
+        elif self.posenc_type == 'dpb':
+            pos_enocoding = self.mlp(positions)
+
+        pos_bias = None if self.complex_transform else self.proj(pos_enocoding)
+
+        return pos_enocoding, pos_bias
+
+class KerP(nn.Module):
+    def __init__(
+        self,
+        feat_dim,
+        sq_factor=4,
+        r=1,
+    ):
+        super().__init__()
+
+        # self.dim = base_dim
+        self.r = r
+        self.pb_encoder1 = PositionEncoder()
+        self.pb_encoder2 = PositionEncoder()
+        self.pb_encoder3 = PositionEncoder()
+        self.r_area = (2 * self.r + 1)**2
+        imnet_in_dim = 32 * 3
+        self.imnets = Siren(in_features=imnet_in_dim, out_features=feat_dim, hidden_features=[feat_dim//sq_factor, feat_dim//sq_factor, feat_dim//sq_factor],
+                                  hidden_layers=2, outermost_linear=True)   #MLP(imnet_in_dim, feat_dim)
+
+    def forward(self, feat, sample_coord, cell):
+        bs = feat.shape[0]
+        r = self.r
+        coord_lr = make_coord(feat.shape[-2:], flatten=False).to(sample_coord.device).permute(2, 0, 1).contiguous(). \
+            unsqueeze(0).expand(bs, 2, *feat.shape[-2:])
+
+        # b, 2, h_in, w_in -> b, 2, h_t, w_t -> b, h_t, w_t, 2  -> b, q, 2
+        sample_coord_k = F.grid_sample(
+            coord_lr, sample_coord.flip(-1), mode='nearest', align_corners=False
+        ).permute(0, 2, 3, 1).contiguous().view(bs, -1, 2).contiguous()
+
+        rel_coord = sample_coord.view(bs, -1, 2).contiguous() - sample_coord_k
+        rel_coord[..., 0] *= feat.shape[-2]
+        rel_coord[..., 1] *= feat.shape[-1]
+
+        dh = torch.linspace(-r, r, 2 * r + 1).to(sample_coord.device)  # * rh
+        dw = torch.linspace(-r, r, 2 * r + 1).to(sample_coord.device)  # * rw
+        # 1, r_area, 2
+        delta = torch.stack(torch.meshgrid(dh, dw, indexing='ij'), axis=-1).view(1, -1, 2).repeat(bs, 1, 1).contiguous()
+
+        # b, 2 -> b, q, 2
+        rel_cell = cell.unsqueeze(1).repeat(1, sample_coord_k.shape[1], 1)
+        rel_cell[..., 0] *= feat.shape[-2]
+        rel_cell[..., 1] *= feat.shape[-1]
+
+        _, pb1 = self.pb_encoder1(rel_cell)  # b, q, 32
+        _, pb2 = self.pb_encoder2(delta)  # b, r_area, 32
+        _, pb3 = self.pb_encoder3(rel_coord)  # b, q, 32
+        # print(pb1.shape, pb2.shape,pb3.shape)
+        pb1 = pb1.unsqueeze(1).repeat(1, self.r_area, 1, 1).view(-1, sample_coord_k.shape[1], 32)
+        pb2 = pb2.unsqueeze(2).repeat(1, 1, sample_coord_k.shape[1], 1).view(-1, sample_coord_k.shape[1],
+                                                                             32)  # b*r_area, q, 32
+        pb3 = pb3.unsqueeze(1).repeat(1, self.r_area, 1, 1).view(-1, sample_coord_k.shape[1], 32)
+        # print(pb1.shape, pb2.shape, pb3.shape)
+        feat_in = torch.cat([pb1, pb2, pb3], dim=-1)
+
+        pred = self.imnets(feat_in)  # b*r_area, q, feat_dim
+        pred = pred.view(bs, self.r_area, sample_coord_k.shape[1], feat.shape[1]).contiguous().permute(0, 3, 1,
+                                                                                                       2).contiguous()
+        kernel = pred.view(bs, feat.shape[1] * self.r_area, -1).contiguous()
+        return kernel
+
+class CUF(nn.Module):
+    def __init__(
+        self,
+        feat_dim,
+        r=1,
+    ):
+        super().__init__()
+
+        self.r = r
+        self.r_area = (2 * self.r + 1)**2
+        self.unfold = torch.nn.Unfold((2 * r + 1, 2 * r + 1), 1, 1)
+        self.imnets = nn.Sequential(
+            nn.Conv2d(feat_dim, feat_dim, 1),
+            nn.LeakyReLU(negative_slope=0.1, inplace=True),
+            nn.Conv2d(feat_dim, 3, 3, 1, 1)
+        )
+
+    def forward(self, feat, kernel, sample_coord, inp):
+        '''
+        feat: B, C, H_in, W_in
+        smaple_coord: B, H_target, W_target, 2
+        cell: B, 2
+        inp:  B, C, H_in, W_in
+        '''
+
+        bs, H_t, W_t, _ = sample_coord.shape
+
+        feat_unf = self.unfold(feat)  # B, C*r_area, H_in*W_in
+        feat_unf = F.grid_sample(feat_unf.view(bs, -1, feat.shape[-2], feat.shape[-1]).contiguous(), sample_coord.flip(-1), mode='bilinear', align_corners=False).view(bs,feat.shape[1]*self.r_area,-1).contiguous()
+
+        pred = feat_unf * kernel  # B, C*r_area, H_target * W_target
+
+        Fold = torch.nn.Fold(sample_coord.shape[1:3], (2 * self.r + 1, 2 * self.r + 1), 1, 1)
+        pred_ = Fold(pred)
+        pred = self.imnets(pred_)
+
+        pred_f = pred + F.grid_sample(inp, sample_coord.flip(-1), mode='bilinear', padding_mode='border', align_corners=False)
+        return pred_f
+

asvr_module/init.py

@@ -0,0 +1,69 @@
+"""
+# --------------------------------------------
+# weights initialization
+# --------------------------------------------
+"""
+from torch.nn import init
+
+"""
+# Kai Zhang, https://github.com/cszn/KAIR
+#
+# Args:
+#   init_type:
+#       normal; normal; xavier_normal; xavier_uniform;
+#       kaiming_normal; kaiming_uniform; orthogonal
+#   init_bn_type:
+#       uniform; constant
+#   gain:
+#       0.2
+"""
+
+def init_fn(m, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2):
+    classname = m.__class__.__name__
+
+    if classname.find('Conv') != -1 or classname.find('Linear') != -1:
+
+        if init_type == 'normal':
+            init.normal_(m.weight.data, 0, 0.1)
+            m.weight.data.clamp_(-1, 1).mul_(gain)
+
+        elif init_type == 'uniform':
+            init.uniform_(m.weight.data, -0.2, 0.2)
+            m.weight.data.mul_(gain)
+
+        elif init_type == 'xavier_normal':
+            init.xavier_normal_(m.weight.data, gain=gain)
+            m.weight.data.clamp_(-1, 1)
+
+        elif init_type == 'xavier_uniform':
+            init.xavier_uniform_(m.weight.data, gain=gain)
+
+        elif init_type == 'kaiming_normal':
+            init.kaiming_normal_(m.weight.data, a=0, mode='fan_in', nonlinearity='relu')
+            m.weight.data.clamp_(-1, 1).mul_(gain)
+
+        elif init_type == 'kaiming_uniform':
+            init.kaiming_uniform_(m.weight.data, a=0, mode='fan_in', nonlinearity='relu')
+            m.weight.data.mul_(gain)
+
+        elif init_type == 'orthogonal':
+            init.orthogonal_(m.weight.data, gain=gain)
+
+        else:
+            raise NotImplementedError('Initialization method [{:s}] is not implemented'.format(init_type))
+
+        if m.bias is not None:
+            m.bias.data.zero_()
+
+    elif classname.find('BatchNorm2d') != -1:
+
+        if init_bn_type == 'uniform':  # preferred
+            if m.affine:
+                init.uniform_(m.weight.data, 0.1, 1.0)
+                init.constant_(m.bias.data, 0.0)
+        elif init_bn_type == 'constant':
+            if m.affine:
+                init.constant_(m.weight.data, 1.0)
+                init.constant_(m.bias.data, 0.0)
+        else:
+            raise NotImplementedError('Initialization method [{:s}] is not implemented'.format(init_bn_type))

asvr_module/network_vidue_refsr_cuf3_siren.py

@@ -0,0 +1,224 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from . import blocks
+from mmcv.ops import ModulatedDeformConv2d, modulated_deform_conv2d
+
+def make_coord(shape):
+    """ Make coordinates at grid centers.
+    """
+    coord_seqs = []
+    for i, n in enumerate(shape):
+        # v0, v1 = -1, 1
+
+        r = 1 / n
+        seq = -1 + r + (2 * r) * torch.arange(n).float()
+        coord_seqs.append(seq)
+    # ret = torch.stack(torch.meshgrid(coord_seqs, indexing='ij'), dim=-1)
+    ret = torch.stack(torch.meshgrid(coord_seqs), dim=-1)
+    return ret
+
+class DynAgg(ModulatedDeformConv2d):
+    '''
+    Use other features to generate offsets and masks.
+    Intialized the offset with precomputed non-local offset.
+    '''
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 padding,
+                 dilation=1,
+                 groups=1,
+                 deform_groups=1,
+                 extra_offset_mask=True):
+        super().__init__(in_channels, out_channels, kernel_size, stride,
+                             padding, dilation, groups, deform_groups)
+        self.extra_offset_mask = extra_offset_mask
+        channels_ = self.deform_groups * 3 * self.kernel_size[
+            0] * self.kernel_size[1]
+        self.conv_offset_mask = nn.Conv2d(
+            self.in_channels,
+            channels_,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            bias=True)
+        self.init_offset()
+
+    def init_offset(self):
+        self.conv_offset_mask.weight.data.zero_()
+        self.conv_offset_mask.bias.data.zero_()
+
+    def forward(self, x, pre_offset):
+        '''
+            Args:
+                pre_offset: precomputed_offset. Size: [b, 2, h, w]
+        '''
+        # pre_offset
+
+        if self.extra_offset_mask:
+            # x = [input, features]
+            out = self.conv_offset_mask(x[1])
+            x = x[0]
+        else:
+            out = self.conv_offset_mask(x)
+        o1, o2, mask = torch.chunk(out, 3, dim=1)
+        offset_ = torch.cat((o1, o2), dim=1)
+
+        offset = offset_ + pre_offset.flip(1).repeat(1, offset_.size(1) // 2, 1, 1)
+
+        # # repeat pre_offset along dim1, shape: [b, 9*groups, h, w, 2]
+        # pre_offset = pre_offset.repeat([1, self.deform_groups, 1, 1, 1])
+        # # the order of offset is [y, x, y, x, ..., y, x]
+        # pre_offset_reorder = torch.zeros_like(offset)
+        # # add pre_offset on y-axis
+        # pre_offset_reorder[:, 0::2, :, :] = pre_offset[:, :, :, :, 1]
+        # # add pre_offset on x-axis
+        # pre_offset_reorder[:, 1::2, :, :] = pre_offset[:, :, :, :, 0]
+        # offset = offset + pre_offset_reorder
+        # print(offset.size())
+        mask = torch.sigmoid(mask)
+
+        offset_mean = torch.mean(torch.abs(offset_))
+        if offset_mean > 100:
+            print('Offset mean is {}, larger than 100.'.format(offset_mean))
+        return modulated_deform_conv2d(x, offset, mask, self.weight, self.bias,
+                           self.stride, self.padding, self.dilation,
+                           self.groups, self.deform_groups)
+
+
+class MRAPAFusion(nn.Module):
+    """Temporal Spatial Attention (TSA) fusion module. It is used in EDVRNet.
+    Args:
+        nf (int): Number of the channels of middle features.
+            Default: 64.
+        ref_nf (int): Number of the channels of middle features.
+            Default: 256.
+    """
+
+    def __init__(self,
+                 nf=64,
+                 ref_nf=256):
+        super().__init__()
+
+        # multi-ref attention (before fusion conv)
+        self.patch_size = 3
+        channels = ref_nf
+        self.conv_emb1 = nn.Sequential(
+            nn.Conv2d(nf, channels, 1),
+            nn.PReLU())
+        self.conv_emb2 = nn.Sequential(
+            nn.Conv2d(ref_nf, channels,
+                      self.patch_size, 1, self.patch_size // 2),
+            nn.PReLU())
+        self.conv_ass = nn.Conv2d(ref_nf, channels * 2,
+                                  self.patch_size, 1, self.patch_size // 2)
+        self.scale = channels ** -0.5
+        self.feat_fusion = nn.Conv2d(
+            nf + channels * 2, nf, 1)
+
+        # spatial attention (after fusion conv)
+        self.spatial_attn = nn.Conv2d(
+            nf + channels * 2, channels * 2, 1)
+        self.spatial_attn_mul1 = nn.Conv2d(
+            channels * 2, channels * 2, 3, padding=1)
+        self.spatial_attn_mul2 = nn.Conv2d(
+            channels * 2, channels * 2, 3, padding=1)
+        self.spatial_attn_add1 = nn.Conv2d(
+            channels * 2, channels * 2, 3, padding=1)
+        self.spatial_attn_add2 = nn.Conv2d(
+            channels * 2, channels * 2, 3, padding=1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+
+    def spatial_padding(self, feats):
+        _, _, h, w = feats.size()
+        pad_h = (4 - h % 4) % 4
+        pad_w = (4 - w % 4) % 4
+        feats = F.pad(feats, [0, pad_w, 0, pad_h], mode='reflect')
+        return feats
+
+    def forward(self, target, refs):
+        n, _, h_input, w_input = target.size()
+        t = len(refs)
+
+        target = self.spatial_padding(target)
+        refs = torch.stack(refs, dim=1).flatten(0, 1)
+        refs = self.spatial_padding(refs)
+        # multi-ref attention
+        embedding_target = self.conv_emb1(target) * self.scale  # (n, c, h, w)
+        embedding_target = embedding_target.permute(0, 2, 3, 1).contiguous().unsqueeze(3)  # (n, h, w, 1, c)
+        embedding_target = embedding_target.contiguous().flatten(0, 2)  # (n*h*w, 1, c)
+        emb = self.conv_emb2(refs).unflatten(0, (n, t))  # (n, t, c, h, w)
+        emb = emb.permute(0, 3, 4, 2, 1).contiguous()  # (n, h, w, c, t)
+        emb = emb.flatten(0, 2)  # (n*h*w, c, t)
+        ass = self.conv_ass(refs).unflatten(0, (n, t))  # (n, t, c*2, h, w)
+        ass = ass.permute(0, 3, 4, 1, 2).contiguous()  # (n, h, w, t, c*2)
+        ass = ass.flatten(0, 2)  # (n*h*w, t, c*2)
+
+        corr_prob = torch.matmul(embedding_target, emb)  # (n*h*w, 1, t)
+        corr_prob = F.softmax(corr_prob, dim=2)
+        refs = torch.matmul(corr_prob, ass).squeeze(1)  # (n*h*w, c*2)
+        refs = refs.unflatten(0, (n, *target.shape[-2:]))  # (n, h, w, c*2)
+        refs = refs.permute(0, 3, 1, 2).contiguous()  # (n, c*2, h, w)
+
+        del embedding_target, emb, ass, corr_prob
+
+        # spatial attention
+        attn = self.lrelu(self.spatial_attn(torch.cat([target, refs], dim=1)))
+        attn_mul = self.spatial_attn_mul2(self.lrelu(self.spatial_attn_mul1(attn)))
+        attn_add = self.spatial_attn_add2(self.lrelu(self.spatial_attn_add1(attn)))
+        attn_mul = torch.sigmoid(attn_mul)
+
+        # after initialization, * 2 makes (attn_mul * 2) to be close to 1.
+        refs = refs * attn_mul * 2 + attn_add
+
+        # fusion
+        feat = self.lrelu(self.feat_fusion(torch.cat([target, refs], dim=1)))
+        return feat[:, :, :h_input, :w_input]
+
+class SynBlock(nn.Module):
+    def __init__(self, nf, ks, groups=8):
+        super(SynBlock, self).__init__()
+        self.xN = 5
+        self.offset_conv1 = nn.Conv2d(
+            nf*2, nf, 3, 1, 1, bias=True)  # concat for diff
+        self.offset_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.dyn_agg = DynAgg(
+            nf,
+            nf,
+            3,
+            stride=1,
+            padding=1,
+            dilation=1,
+            deform_groups=groups,
+            extra_offset_mask=True)
+        self.head = MRAPAFusion(nf=nf, ref_nf=nf)
+        self.body = blocks.ResBlock(nf, nf, kernel_size=ks, stride=1, se=False)
+        self.tail = nn.Sequential(
+            nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1))
+        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+
+    def forward(self, x, pre_offset_list, img_ref_feat_list, size):  #fea, flow_list, vec, out_size):  #
+
+        swapped_feat_list = []
+        for pre_offset, img_ref_feat in zip(pre_offset_list, img_ref_feat_list):
+            _, _, H, W = pre_offset.shape
+            TH, TW = size
+            scale = int(TW / W)
+            # print(scale, TW, W)
+            pre_offset = F.interpolate(pre_offset, size=size, mode='bilinear') * scale
+            offset = torch.cat([x, img_ref_feat], 1)
+            offset = self.lrelu(self.offset_conv1(offset))
+            offset = self.lrelu(self.offset_conv2(offset))
+            swapped_feat = self.lrelu(
+                self.dyn_agg([img_ref_feat, offset], pre_offset))
+            swapped_feat_list.append(swapped_feat)
+
+        h = self.head(x, swapped_feat_list)
+        h = self.body(h) + x
+        x = self.tail(h)
+        return x

asvr_module/pytorch_pwc/correlation/correlation.py

@@ -0,0 +1,435 @@
+#!/usr/bin/env python
+
+import torch
+
+import cupy
+import re
+
+kernel_Correlation_rearrange = '''
+	extern "C" __global__ void kernel_Correlation_rearrange(
+		const int n,
+		const float* input,
+		float* output
+	) {
+	  int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x;
+
+	  if (intIndex >= n) {
+	    return;
+	  }
+
+	  int intSample = blockIdx.z;
+	  int intChannel = blockIdx.y;
+
+	  float fltValue = input[(((intSample * SIZE_1(input)) + intChannel) * SIZE_2(input) * SIZE_3(input)) + intIndex];
+
+	  __syncthreads();
+
+	  int intPaddedY = (intIndex / SIZE_3(input)) + 4;
+	  int intPaddedX = (intIndex % SIZE_3(input)) + 4;
+	  int intRearrange = ((SIZE_3(input) + 8) * intPaddedY) + intPaddedX;
+
+	  output[(((intSample * SIZE_1(output) * SIZE_2(output)) + intRearrange) * SIZE_1(input)) + intChannel] = fltValue;
+	}
+'''
+
+kernel_Correlation_updateOutput = '''
+	extern "C" __global__ void kernel_Correlation_updateOutput(
+	  const int n,
+	  const float* rbot0,
+	  const float* rbot1,
+	  float* top
+	) {
+	  extern __shared__ char patch_data_char[];
+
+	  float *patch_data = (float *)patch_data_char;
+
+	  // First (upper left) position of kernel upper-left corner in current center position of neighborhood in image 1
+	  int x1 = blockIdx.x + 4;
+	  int y1 = blockIdx.y + 4;
+	  int item = blockIdx.z;
+	  int ch_off = threadIdx.x;
+
+	  // Load 3D patch into shared shared memory
+	  for (int j = 0; j < 1; j++) { // HEIGHT
+	    for (int i = 0; i < 1; i++) { // WIDTH
+	      int ji_off = (j + i) * SIZE_3(rbot0);
+	      for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS
+	        int idx1 = ((item * SIZE_1(rbot0) + y1+j) * SIZE_2(rbot0) + x1+i) * SIZE_3(rbot0) + ch;
+	        int idxPatchData = ji_off + ch;
+	        patch_data[idxPatchData] = rbot0[idx1];
+	      }
+	    }
+	  }
+
+	  __syncthreads();
+
+	  __shared__ float sum[32];
+
+	  // Compute correlation
+	  for (int top_channel = 0; top_channel < SIZE_1(top); top_channel++) {
+	    sum[ch_off] = 0;
+
+	    int s2o = top_channel % 9 - 4;
+	    int s2p = top_channel / 9 - 4;
+
+	    for (int j = 0; j < 1; j++) { // HEIGHT
+	      for (int i = 0; i < 1; i++) { // WIDTH
+	        int ji_off = (j + i) * SIZE_3(rbot0);
+	        for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS
+	          int x2 = x1 + s2o;
+	          int y2 = y1 + s2p;
+
+	          int idxPatchData = ji_off + ch;
+	          int idx2 = ((item * SIZE_1(rbot0) + y2+j) * SIZE_2(rbot0) + x2+i) * SIZE_3(rbot0) + ch;
+
+	          sum[ch_off] += patch_data[idxPatchData] * rbot1[idx2];
+	        }
+	      }
+	    }
+
+	    __syncthreads();
+
+	    if (ch_off == 0) {
+	      float total_sum = 0;
+	      for (int idx = 0; idx < 32; idx++) {
+	        total_sum += sum[idx];
+	      }
+	      const int sumelems = SIZE_3(rbot0);
+	      const int index = ((top_channel*SIZE_2(top) + blockIdx.y)*SIZE_3(top))+blockIdx.x;
+	      top[index + item*SIZE_1(top)*SIZE_2(top)*SIZE_3(top)] = total_sum / (float)sumelems;
+	    }
+	  }
+	}
+'''
+
+kernel_Correlation_updateGradFirst = '''
+	#define ROUND_OFF 50000
+
+	extern "C" __global__ void kernel_Correlation_updateGradFirst(
+	  const int n,
+	  const int intSample,
+	  const float* rbot0,
+	  const float* rbot1,
+	  const float* gradOutput,
+	  float* gradFirst,
+	  float* gradSecond
+	) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) {
+	  int n = intIndex % SIZE_1(gradFirst); // channels
+	  int l = (intIndex / SIZE_1(gradFirst)) % SIZE_3(gradFirst) + 4; // w-pos
+	  int m = (intIndex / SIZE_1(gradFirst) / SIZE_3(gradFirst)) % SIZE_2(gradFirst) + 4; // h-pos
+
+	  // round_off is a trick to enable integer division with ceil, even for negative numbers
+	  // We use a large offset, for the inner part not to become negative.
+	  const int round_off = ROUND_OFF;
+	  const int round_off_s1 = round_off;
+
+	  // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior:
+	  int xmin = (l - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4)
+	  int ymin = (m - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4)
+
+	  // Same here:
+	  int xmax = (l - 4 + round_off_s1) - round_off; // floor (l - 4)
+	  int ymax = (m - 4 + round_off_s1) - round_off; // floor (m - 4)
+
+	  float sum = 0;
+	  if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) {
+	    xmin = max(0,xmin);
+	    xmax = min(SIZE_3(gradOutput)-1,xmax);
+
+	    ymin = max(0,ymin);
+	    ymax = min(SIZE_2(gradOutput)-1,ymax);
+
+	    for (int p = -4; p <= 4; p++) {
+	      for (int o = -4; o <= 4; o++) {
+	        // Get rbot1 data:
+	        int s2o = o;
+	        int s2p = p;
+	        int idxbot1 = ((intSample * SIZE_1(rbot0) + (m+s2p)) * SIZE_2(rbot0) + (l+s2o)) * SIZE_3(rbot0) + n;
+	        float bot1tmp = rbot1[idxbot1]; // rbot1[l+s2o,m+s2p,n]
+
+	        // Index offset for gradOutput in following loops:
+	        int op = (p+4) * 9 + (o+4); // index[o,p]
+	        int idxopoffset = (intSample * SIZE_1(gradOutput) + op);
+
+	        for (int y = ymin; y <= ymax; y++) {
+	          for (int x = xmin; x <= xmax; x++) {
+	            int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p]
+	            sum += gradOutput[idxgradOutput] * bot1tmp;
+	          }
+	        }
+	      }
+	    }
+	  }
+	  const int sumelems = SIZE_1(gradFirst);
+	  const int bot0index = ((n * SIZE_2(gradFirst)) + (m-4)) * SIZE_3(gradFirst) + (l-4);
+	  gradFirst[bot0index + intSample*SIZE_1(gradFirst)*SIZE_2(gradFirst)*SIZE_3(gradFirst)] = sum / (float)sumelems;
+	} }
+'''
+
+kernel_Correlation_updateGradSecond = '''
+	#define ROUND_OFF 50000
+
+	extern "C" __global__ void kernel_Correlation_updateGradSecond(
+	  const int n,
+	  const int intSample,
+	  const float* rbot0,
+	  const float* rbot1,
+	  const float* gradOutput,
+	  float* gradFirst,
+	  float* gradSecond
+	) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) {
+	  int n = intIndex % SIZE_1(gradSecond); // channels
+	  int l = (intIndex / SIZE_1(gradSecond)) % SIZE_3(gradSecond) + 4; // w-pos
+	  int m = (intIndex / SIZE_1(gradSecond) / SIZE_3(gradSecond)) % SIZE_2(gradSecond) + 4; // h-pos
+
+	  // round_off is a trick to enable integer division with ceil, even for negative numbers
+	  // We use a large offset, for the inner part not to become negative.
+	  const int round_off = ROUND_OFF;
+	  const int round_off_s1 = round_off;
+
+	  float sum = 0;
+	  for (int p = -4; p <= 4; p++) {
+	    for (int o = -4; o <= 4; o++) {
+	      int s2o = o;
+	      int s2p = p;
+
+	      //Get X,Y ranges and clamp
+	      // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior:
+	      int xmin = (l - 4 - s2o + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o)
+	      int ymin = (m - 4 - s2p + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o)
+
+	      // Same here:
+	      int xmax = (l - 4 - s2o + round_off_s1) - round_off; // floor (l - 4 - s2o)
+	      int ymax = (m - 4 - s2p + round_off_s1) - round_off; // floor (m - 4 - s2p)
+
+	      if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) {
+	        xmin = max(0,xmin);
+	        xmax = min(SIZE_3(gradOutput)-1,xmax);
+
+	        ymin = max(0,ymin);
+	        ymax = min(SIZE_2(gradOutput)-1,ymax);
+
+	        // Get rbot0 data:
+	        int idxbot0 = ((intSample * SIZE_1(rbot0) + (m-s2p)) * SIZE_2(rbot0) + (l-s2o)) * SIZE_3(rbot0) + n;
+	        float bot0tmp = rbot0[idxbot0]; // rbot1[l+s2o,m+s2p,n]
+
+	        // Index offset for gradOutput in following loops:
+	        int op = (p+4) * 9 + (o+4); // index[o,p]
+	        int idxopoffset = (intSample * SIZE_1(gradOutput) + op);
+
+	        for (int y = ymin; y <= ymax; y++) {
+	          for (int x = xmin; x <= xmax; x++) {
+	            int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p]
+	            sum += gradOutput[idxgradOutput] * bot0tmp;
+	          }
+	        }
+	      }
+	    }
+	  }
+	  const int sumelems = SIZE_1(gradSecond);
+	  const int bot1index = ((n * SIZE_2(gradSecond)) + (m-4)) * SIZE_3(gradSecond) + (l-4);
+	  gradSecond[bot1index + intSample*SIZE_1(gradSecond)*SIZE_2(gradSecond)*SIZE_3(gradSecond)] = sum / (float)sumelems;
+	} }
+'''
+
+
+def cupy_kernel(strFunction, objVariables):
+    strKernel = globals()[strFunction]
+
+    while True:
+        objMatch = re.search('(SIZE_)([0-4])(\()([^\)]*)(\))', strKernel)
+
+        if objMatch is None:
+            break
+        # end
+
+        intArg = int(objMatch.group(2))
+
+        strTensor = objMatch.group(4)
+        intSizes = objVariables[strTensor].size()
+
+        strKernel = strKernel.replace(objMatch.group(), str(intSizes[intArg]))
+    # end
+
+    while True:
+        objMatch = re.search('(VALUE_)([0-4])(\()([^\)]+)(\))', strKernel)
+
+        if objMatch is None:
+            break
+        # end
+
+        intArgs = int(objMatch.group(2))
+        strArgs = objMatch.group(4).split(',')
+
+        strTensor = strArgs[0]
+        intStrides = objVariables[strTensor].stride()
+        strIndex = ['((' + strArgs[intArg + 1].replace('{', '(').replace('}', ')').strip() + ')*' + str(
+            intStrides[intArg]) + ')' for intArg in range(intArgs)]
+
+        strKernel = strKernel.replace(objMatch.group(0), strTensor + '[' + str.join('+', strIndex) + ']')
+    # end
+
+    return strKernel
+
+
+# end
+
+# @cupy.memoize(for_each_device=True)
+# def cupy_launch(strFunction, strKernel):
+#     return cupy.cuda.compile_with_cache(strKernel).get_function(strFunction)
+
+compiled_modules = {}
+
+def get_device_id():
+    return cupy.cuda.Device().id
+
+def get_or_compile_module(device_id, strKernel):
+    if device_id not in compiled_modules:
+        compiled_modules[device_id] = {}
+
+    if strKernel not in compiled_modules[device_id]:
+        module = cupy.RawModule(code=strKernel)
+        compiled_modules[device_id][strKernel] = module
+
+    return compiled_modules[device_id][strKernel]
+
+def cupy_launch(strFunction, strKernel):
+    device_id = get_device_id()
+    module = get_or_compile_module(device_id, strKernel)
+    function = module.get_function(strFunction)
+    return function
+
+# end
+
+class _FunctionCorrelation(torch.autograd.Function):
+    @staticmethod
+    def forward(self, first, second):
+        rbot0 = first.new_zeros([first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1]])
+        rbot1 = first.new_zeros([first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1]])
+
+        self.save_for_backward(first, second, rbot0, rbot1)
+
+        assert (first.is_contiguous() == True)
+        assert (second.is_contiguous() == True)
+
+        output = first.new_zeros([first.shape[0], 81, first.shape[2], first.shape[3]])
+
+        if first.is_cuda == True:
+            n = first.shape[2] * first.shape[3]
+            cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
+                'input': first,
+                'output': rbot0
+            }))(
+                grid=tuple([int((n + 16 - 1) / 16), first.shape[1], first.shape[0]]),
+                block=tuple([16, 1, 1]),
+                args=[n, first.data_ptr(), rbot0.data_ptr()]
+            )
+
+            n = second.shape[2] * second.shape[3]
+            cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
+                'input': second,
+                'output': rbot1
+            }))(
+                grid=tuple([int((n + 16 - 1) / 16), second.shape[1], second.shape[0]]),
+                block=tuple([16, 1, 1]),
+                args=[n, second.data_ptr(), rbot1.data_ptr()]
+            )
+
+            n = output.shape[1] * output.shape[2] * output.shape[3]
+            cupy_launch('kernel_Correlation_updateOutput', cupy_kernel('kernel_Correlation_updateOutput', {
+                'rbot0': rbot0,
+                'rbot1': rbot1,
+                'top': output
+            }))(
+                grid=tuple([output.shape[3], output.shape[2], output.shape[0]]),
+                block=tuple([32, 1, 1]),
+                shared_mem=first.shape[1] * 4,
+                args=[n, rbot0.data_ptr(), rbot1.data_ptr(), output.data_ptr()]
+            )
+
+        elif first.is_cuda == False:
+            raise NotImplementedError()
+
+        # end
+
+        return output
+
+    # end
+
+    @staticmethod
+    def backward(self, gradOutput):
+        first, second, rbot0, rbot1 = self.saved_tensors
+
+        assert (gradOutput.is_contiguous() == True)
+
+        gradFirst = first.new_zeros([first.shape[0], first.shape[1], first.shape[2], first.shape[3]]) if \
+        self.needs_input_grad[0] == True else None
+        gradSecond = first.new_zeros([first.shape[0], first.shape[1], first.shape[2], first.shape[3]]) if \
+        self.needs_input_grad[1] == True else None
+
+        if first.is_cuda == True:
+            if gradFirst is not None:
+                for intSample in range(first.shape[0]):
+                    n = first.shape[1] * first.shape[2] * first.shape[3]
+                    cupy_launch('kernel_Correlation_updateGradFirst',
+                                cupy_kernel('kernel_Correlation_updateGradFirst', {
+                                    'rbot0': rbot0,
+                                    'rbot1': rbot1,
+                                    'gradOutput': gradOutput,
+                                    'gradFirst': gradFirst,
+                                    'gradSecond': None
+                                }))(
+                        grid=tuple([int((n + 512 - 1) / 512), 1, 1]),
+                        block=tuple([512, 1, 1]),
+                        args=[n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(),
+                              gradFirst.data_ptr(), None]
+                    )
+            # end
+            # end
+
+            if gradSecond is not None:
+                for intSample in range(first.shape[0]):
+                    n = first.shape[1] * first.shape[2] * first.shape[3]
+                    cupy_launch('kernel_Correlation_updateGradSecond',
+                                cupy_kernel('kernel_Correlation_updateGradSecond', {
+                                    'rbot0': rbot0,
+                                    'rbot1': rbot1,
+                                    'gradOutput': gradOutput,
+                                    'gradFirst': None,
+                                    'gradSecond': gradSecond
+                                }))(
+                        grid=tuple([int((n + 512 - 1) / 512), 1, 1]),
+                        block=tuple([512, 1, 1]),
+                        args=[n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), None,
+                              gradSecond.data_ptr()]
+                    )
+            # end
+        # end
+
+        elif first.is_cuda == False:
+            raise NotImplementedError()
+
+        # end
+
+        return gradFirst, gradSecond
+
+
+# end
+# end
+
+def FunctionCorrelation(tenFirst, tenSecond):
+    return _FunctionCorrelation.apply(tenFirst, tenSecond)
+
+
+# end
+
+class ModuleCorrelation(torch.nn.Module):
+    def __init__(self):
+        super(ModuleCorrelation, self).__init__()
+
+    # end
+
+    def forward(self, tenFirst, tenSecond):
+        return _FunctionCorrelation.apply(tenFirst, tenSecond)
+# end
+# end

asvr_module/pytorch_pwc/extract_flow.py

@@ -0,0 +1,32 @@
+import math
+import torch
+
+# im1_torch, im2_torch in shape (N, C, H, W)
+def extract_flow_torch(model, im1_torch, im2_torch):
+    # interpolate image, make new_H, mew_W divide by 64
+    assert im1_torch.shape == im2_torch.shape
+    N, C, H, W = im1_torch.shape
+    device = im1_torch.device
+    new_H = int(math.floor(math.ceil(H / 64.0) * 64.0))
+    new_W = int(math.floor(math.ceil(W / 64.0) * 64.0))
+    im1_torch = torch.nn.functional.interpolate(input=im1_torch, size=(new_H, new_W), mode='bilinear',
+                                                 align_corners=False)
+    im2_torch = torch.nn.functional.interpolate(input=im2_torch, size=(new_H, new_W), mode='bilinear',
+                                                 align_corners=False)
+    model.eval()
+    with torch.no_grad():
+        flo12 = model(im1_torch, im2_torch)
+    flo12 = 20.0 * torch.nn.functional.interpolate(input=flo12, size=(H, W), mode='bilinear',
+                                                          align_corners=False)
+    flo12[:, 0, :, :] *= float(W) / float(new_W)
+    flo12[:, 1, :, :] *= float(H) / float(new_H)
+    return flo12
+
+# im1_np, im2_np in shape (C, H, W)
+def extract_flow_np(model, im1_np, im2_np):
+    im1_torch = torch.from_numpy(im1_np).unsqueeze(0).to(torch.device('cuda'))
+    im2_torch = torch.from_numpy(im2_np).unsqueeze(0).to(torch.device('cuda'))
+    flo12_torch = extract_flow_torch(model, im1_torch, im2_torch)
+    flo12_np = flo12_torch.detach().cpu().squeeze(0).numpy()
+    return flo12_np
+

asvr_module/pytorch_pwc/pwc.py

@@ -0,0 +1,242 @@
+from .correlation import correlation # the custom cost volume layer
+import torch
+
+##########################################################
+
+assert(int(str('').join(torch.__version__.split('.')[0:2])) >= 13) # requires at least pytorch version 1.3.0
+
+# torch.set_grad_enabled(False)
+
+torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance
+
+arguments_strModel = 'default'
+
+backwarp_tenGrid = {}
+backwarp_tenPartial = {}
+
+def backwarp(tenInput, tenFlow):
+    if (str(tenFlow.shape)+str(tenFlow.device)) not in backwarp_tenGrid:
+        tenHor = torch.linspace(-1.0 + (1.0 / tenFlow.shape[3]), 1.0 - (1.0 / tenFlow.shape[3]), tenFlow.shape[3]).view(1, 1, 1, -1).expand(-1, -1, tenFlow.shape[2], -1)
+        tenVer = torch.linspace(-1.0 + (1.0 / tenFlow.shape[2]), 1.0 - (1.0 / tenFlow.shape[2]), tenFlow.shape[2]).view(1, 1, -1, 1).expand(-1, -1, -1, tenFlow.shape[3])
+
+        backwarp_tenGrid[str(tenFlow.shape) + str(tenFlow.device)] = torch.cat([ tenHor, tenVer ], 1).to(tenFlow.device)
+    # end
+
+    if (str(tenFlow.shape)+str(tenFlow.device)) not in backwarp_tenPartial:
+        backwarp_tenPartial[str(tenFlow.shape)+str(tenFlow.device)] = tenFlow.new_ones([ tenFlow.shape[0], 1, tenFlow.shape[2], tenFlow.shape[3] ])
+    # end
+
+    tenFlow = torch.cat([ tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0) ], 1)
+    tenInput = torch.cat([ tenInput, backwarp_tenPartial[str(tenFlow.shape)+str(tenFlow.device)] ], 1)
+
+    tenOutput = torch.nn.functional.grid_sample(input=tenInput, grid=(backwarp_tenGrid[str(tenFlow.shape) + str(tenFlow.device)] + tenFlow).permute(0, 2, 3, 1), mode='bilinear', padding_mode='zeros', align_corners=False)
+
+    tenMask = tenOutput[:, -1:, :, :]; tenMask[tenMask > 0.999] = 1.0; tenMask[tenMask < 1.0] = 0.0
+
+    return tenOutput[:, :-1, :, :] * tenMask
+# end
+
+##########################################################
+
+class PWCNet(torch.nn.Module):
+    def __init__(self):
+        super(PWCNet, self).__init__()
+
+        class Extractor(torch.nn.Module):
+            def __init__(self):
+                super(Extractor, self).__init__()
+
+                self.netOne = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+
+                self.netTwo = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=2, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+
+                self.netThr = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+
+                self.netFou = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=64, out_channels=96, kernel_size=3, stride=2, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+
+                self.netFiv = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=96, out_channels=128, kernel_size=3, stride=2, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+
+                self.netSix = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=128, out_channels=196, kernel_size=3, stride=2, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+            # end
+
+            def forward(self, tenInput):
+                tenOne = self.netOne(tenInput)
+                tenTwo = self.netTwo(tenOne)
+                tenThr = self.netThr(tenTwo)
+                tenFou = self.netFou(tenThr)
+                tenFiv = self.netFiv(tenFou)
+                tenSix = self.netSix(tenFiv)
+
+                return [ tenOne, tenTwo, tenThr, tenFou, tenFiv, tenSix ]
+            # end
+        # end
+
+        class Decoder(torch.nn.Module):
+            def __init__(self, intLevel):
+                super(Decoder, self).__init__()
+
+                intPrevious = [ None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None ][intLevel + 1]
+                intCurrent = [ None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None ][intLevel + 0]
+
+                if intLevel < 6: self.netUpflow = torch.nn.ConvTranspose2d(in_channels=2, out_channels=2, kernel_size=4, stride=2, padding=1)
+                if intLevel < 6: self.netUpfeat = torch.nn.ConvTranspose2d(in_channels=intPrevious + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=4, stride=2, padding=1)
+                if intLevel < 6: self.fltBackwarp = [ None, None, None, 5.0, 2.5, 1.25, 0.625, None ][intLevel + 1]
+
+                self.netOne = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=intCurrent, out_channels=128, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+
+                self.netTwo = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=intCurrent + 128, out_channels=128, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+
+                self.netThr = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=intCurrent + 128 + 128, out_channels=96, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+
+                self.netFou = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96, out_channels=64, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+
+                self.netFiv = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64, out_channels=32, kernel_size=3, stride=1, padding=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+                )
+
+                self.netSix = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=3, stride=1, padding=1)
+                )
+            # end
+
+            def forward(self, tenFirst, tenSecond, objPrevious):
+                tenFlow = None
+                tenFeat = None
+
+                if objPrevious is None:
+                    tenFlow = None
+                    tenFeat = None
+                    tenVolume = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tenFirst=tenFirst, tenSecond=tenSecond), negative_slope=0.1, inplace=False)
+
+                    tenFeat = torch.cat([ tenVolume ], 1)
+
+                elif objPrevious is not None:
+                    tenFlow = self.netUpflow(objPrevious['tenFlow'])
+                    tenFeat = self.netUpfeat(objPrevious['tenFeat'])
+
+                    tenVolume = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tenFirst=tenFirst, tenSecond=backwarp(tenInput=tenSecond, tenFlow=tenFlow * self.fltBackwarp)), negative_slope=0.1, inplace=False)
+
+                    tenFeat = torch.cat([ tenVolume, tenFirst, tenFlow, tenFeat ], 1)
+
+                # end
+
+                tenFeat = torch.cat([ self.netOne(tenFeat), tenFeat ], 1)
+                tenFeat = torch.cat([ self.netTwo(tenFeat), tenFeat ], 1)
+                tenFeat = torch.cat([ self.netThr(tenFeat), tenFeat ], 1)
+                tenFeat = torch.cat([ self.netFou(tenFeat), tenFeat ], 1)
+                tenFeat = torch.cat([ self.netFiv(tenFeat), tenFeat ], 1)
+
+                tenFlow = self.netSix(tenFeat)
+
+                return {
+                    'tenFlow': tenFlow,
+                    'tenFeat': tenFeat
+                }
+            # end
+        # end
+
+        class Refiner(torch.nn.Module):
+            def __init__(self):
+                super(Refiner, self).__init__()
+
+                self.netMain = torch.nn.Sequential(
+                    torch.nn.Conv2d(in_channels=81 + 32 + 2 + 2 + 128 + 128 + 96 + 64 + 32, out_channels=128, kernel_size=3, stride=1, padding=1, dilation=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=2, dilation=2),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=4, dilation=4),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=128, out_channels=96, kernel_size=3, stride=1, padding=8, dilation=8),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=96, out_channels=64, kernel_size=3, stride=1, padding=16, dilation=16),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1, dilation=1),
+                    torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+                    torch.nn.Conv2d(in_channels=32, out_channels=2, kernel_size=3, stride=1, padding=1, dilation=1)
+                )
+            # end
+
+            def forward(self, tenInput):
+                return self.netMain(tenInput)
+            # end
+        # end
+
+        self.netExtractor = Extractor()
+
+        self.netTwo = Decoder(2)
+        self.netThr = Decoder(3)
+        self.netFou = Decoder(4)
+        self.netFiv = Decoder(5)
+        self.netSix = Decoder(6)
+
+        self.netRefiner = Refiner()
+
+        self.load_state_dict({ strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.hub.load_state_dict_from_url(url='http://content.sniklaus.com/github/pytorch-pwc/network-' + arguments_strModel + '.pytorch', file_name='pwc-' + arguments_strModel).items() })
+    # end
+
+    def forward(self, tenFirst, tenSecond):
+        tenFirst = self.netExtractor(tenFirst)
+        tenSecond = self.netExtractor(tenSecond)
+
+        objEstimate = self.netSix(tenFirst[-1], tenSecond[-1], None)
+        objEstimate = self.netFiv(tenFirst[-2], tenSecond[-2], objEstimate)
+        objEstimate = self.netFou(tenFirst[-3], tenSecond[-3], objEstimate)
+        objEstimate = self.netThr(tenFirst[-4], tenSecond[-4], objEstimate)
+        objEstimate = self.netTwo(tenFirst[-5], tenSecond[-5], objEstimate)
+
+        return objEstimate['tenFlow'] + self.netRefiner(objEstimate['tenFeat'])

asvr_module/refsrrnn_adists_fgda_only_future.py

@@ -0,0 +1,214 @@
+from .components import ResBlocks, D
+from .pytorch_pwc.extract_flow import extract_flow_torch
+from .pytorch_pwc.pwc import PWCNet
+import torch
+import torch.nn as nn
+from .warp import warp
+import torch.nn.functional as F
+from .cuf_siren import KerP, CUF
+# # from models.refsrrnn_fgda import SynBlock
+from .network_vidue_refsr_cuf3_siren import SynBlock
+from .ADISTS_allstage import ADISTS
+
+def std(img, window_size=3):
+    assert window_size % 2 == 1
+    pad = window_size // 2
+
+    # calculate std on the mean image of the color channels
+    # img = torch.mean(img, dim=1, keepdim=True)
+    N, C, H, W = img.shape
+    img = nn.functional.pad(img, [pad] * 4, mode='reflect')
+    img = nn.functional.unfold(img, kernel_size=window_size)
+    img = img.view(N, C, window_size * window_size, H, W)
+    img = img - torch.mean(img, dim=2, keepdim=True)
+    # print(img.shape)
+    img = img * img
+    img = torch.mean(img, dim=2, keepdim=True)
+    # print(img.shape)
+    img = torch.sqrt(img)
+    # print(img.shape)
+    img = img.squeeze(2)
+    return img
+
+def generate_alpha(input, lower=1, upper=5):
+    N, C, H, W = input.shape
+    ratio = input.new_ones((N, C, H, W)) * 0.5
+    input_std = std(input)
+    ratio[input_std < lower] = torch.sigmoid((input_std - lower))[input_std < lower]
+    ratio[input_std > upper] = torch.sigmoid((input_std - upper))[input_std > upper]
+    ratio = ratio.detach()
+
+    return ratio
+
+class RefsrRNN(nn.Module):
+    def __init__(self, img_channels=3, num_resblocks=15, num_channels=64, count=2):
+        super(RefsrRNN, self).__init__()
+        self.num_channels = num_channels
+        self.pwcnet = PWCNet()
+        self.forward_rnn = ResBlocks(input_channels=img_channels + num_channels*2, num_resblocks=num_resblocks, num_channels=num_channels)
+        self.forward_rnn2 = ResBlocks(input_channels=img_channels + num_channels*2, num_resblocks=num_resblocks, num_channels=num_channels)
+        self.d = D(in_channels=num_channels * 2, mid_channels=num_channels * 2, out_channels=num_channels)
+        self.kernel_predict = KerP(num_channels)
+        self.upsample = CUF(num_channels)
+        self.predict = SynBlock(num_channels, ks=3)
+        self.count = count
+        self.adists = ADISTS()
+        self.multi_prior_fusion = nn.Conv2d(64+128+256+512+512,num_channels,1,1,0)
+        for param in self.pwcnet.parameters():
+            param.requires_grad = False
+        for param in self.adists.parameters():
+            param.requires_grad = False
+
+    # def trainable_parameters(self):
+    #     return [{'params':self.forward_rnn.parameters()}, {'params':self.d.parameters()}, {'params':self.kernel_predict.parameters()}, {'params':self.upsample.parameters()}]
+
+    def forward(self, images, scale, coord, cell, isTrain=True):
+        seqn_not_pad = torch.stack(images, dim=1)
+        N, T, C, H, W = seqn_not_pad.shape
+        seqn_not_pad_ = F.interpolate(seqn_not_pad.view(-1, C, H, W).contiguous(), size=(coord.shape[-3:-1]), mode='bicubic')
+        seqn_not_pad_ = seqn_not_pad_.view(N, T, C, coord.shape[-3], coord.shape[-2]).contiguous()
+        seqdn = torch.empty_like(seqn_not_pad_)
+        del seqn_not_pad_
+
+        # border_queue = []
+        # size_h, size_w = int(H * self.border_ratio), int(W * self.border_ratio)
+
+        # reflect pad seqn and noise_level_map
+        seqn = torch.empty((N, T + self.count, C, H, W), device=seqn_not_pad.device)
+        seqn[:, :T] = seqn_not_pad
+        for i in range(self.count):
+            seqn[:, T + i] = seqn_not_pad[:, T - 2 - i]
+
+        # flatten_map = generate_alpha(seqn.reshape(-1, C, H, W).contiguous()*255.).view(N, T+self.count*2, C, H, W).contiguous()
+        # flatten_map = self.adists(seqn.reshape(-1, C, H, W).contiguous()).view(N, T+self.count, self.num_channels, H, W).contiguous()
+
+        inp1, inp2 = seqn[:, 1:].reshape(-1, C, H, W).contiguous(), seqn[:, :-1].reshape(-1, C, H, W).contiguous()
+        flow = extract_flow_torch(self.pwcnet, inp1, inp2)
+        flow_queue = flow.view(N, T + self.count-1, 2, H, W).contiguous()
+
+        hidden_list = []
+        prior_list = []
+        init_forward_h = torch.zeros((N, self.num_channels, H, W), device=seqn.device)
+        flatten_map = self.multi_prior_fusion(self.adists(seqn[:, 0]))
+        h_ = self.forward_rnn(torch.cat((seqn[:, 0], flatten_map, init_forward_h), dim=1))
+        hidden_list.append(h_)
+        prior_list.append(flatten_map)
+        arb_up_kernel = self.kernel_predict(h_, coord, cell)
+        for j in range(1, self.count):
+            h_, _ = warp(h_, flow_queue[:, j-1])
+            flatten_map = self.multi_prior_fusion(self.adists(seqn[:, j]))
+            h_ = self.forward_rnn(torch.cat((seqn[:, j], flatten_map, h_), dim=1))
+            hidden_list.append(h_)
+            prior_list.append(flatten_map)
+
+        for i in range(self.count, T + self.count):
+            # if i == self.count:
+            #     for j in range(i, i+self.count):
+            #         h_, _ = warp(h_, flow_queue[:, j - 1])
+            #         h_ = self.forward_rnn(torch.cat((seqn[:, j], flatten_map[:, j], h_), dim=1))
+            #         hidden_list.append(h_)
+            # else:
+            #     #print(hidden_list[-1].shape)
+            #     #print(flow_queue[:, i + self.count - 1].shape)
+            h_, _ = warp(hidden_list[-1], flow_queue[:, i-1])
+            flatten_map = self.multi_prior_fusion(self.adists(seqn[:, i]))
+            h_ = self.forward_rnn(torch.cat((seqn[:, i], flatten_map, h_), dim=1))
+            hidden_list.append(h_)
+            prior_list.append(flatten_map)
+            if i > self.count:
+                hidden_list.pop(0)
+                prior_list.pop(0)
+
+            assert len(hidden_list) == self.count+1
+
+            refsr_flow_list = []
+            for j in range(1, self.count+1):
+                refsr_flow_list.append(extract_flow_torch(self.pwcnet, seqn[:, i-self.count], seqn[:, i-self.count + j]))
+
+            assert len(refsr_flow_list) == self.count
+
+            hidden = hidden_list[0]
+            fusion_h = self.predict(hidden, refsr_flow_list, hidden_list[1:], hidden.size()[-2:])
+
+            h = self.forward_rnn2(torch.cat((seqn[:, i-self.count], prior_list[0], fusion_h), dim=1))
+            res = self.d(torch.cat([h, hidden], dim=1))
+            seqdn[:, i-self.count] = self.upsample(res, arb_up_kernel, coord, seqn[:, i-self.count])
+            
+
+        return seqdn
+
+    def test_forward(self, images, arb_up_kernel, coord):
+        seqn_not_pad = torch.stack(images, dim=1)
+        print(f"{seqn_not_pad.shape=}")
+        N, T, C, H, W = seqn_not_pad.shape
+        seqn_not_pad_ = F.interpolate(seqn_not_pad.view(-1, C, H, W).contiguous(), size=(coord.shape[-3:-1]), mode='bicubic')
+        seqn_not_pad_ = seqn_not_pad_.view(N, T, C, coord.shape[-3], coord.shape[-2]).contiguous()
+        seqdn = torch.empty_like(seqn_not_pad_)
+        del seqn_not_pad_
+
+        # border_queue = []
+        # size_h, size_w = int(H * self.border_ratio), int(W * self.border_ratio)
+
+        # reflect pad seqn and noise_level_map
+        seqn = torch.empty((N, T + self.count, C, H, W), device=seqn_not_pad.device)
+        seqn[:, :T] = seqn_not_pad
+        for i in range(self.count):
+            seqn[:, T + i] = seqn_not_pad[:, T - 2 - i]
+
+        # flatten_map = generate_alpha(seqn.reshape(-1, C, H, W).contiguous()*255.).view(N, T+self.count*2, C, H, W).contiguous()
+        # flatten_map = self.adists(seqn.reshape(-1, C, H, W).contiguous()).view(N, T+self.count, self.num_channels, H, W).contiguous()
+
+        inp1, inp2 = seqn[:, 1:].reshape(-1, C, H, W).contiguous(), seqn[:, :-1].reshape(-1, C, H, W).contiguous()
+        flow = extract_flow_torch(self.pwcnet, inp1, inp2)
+        flow_queue = flow.view(N, T + self.count-1, 2, H, W).contiguous()
+
+        hidden_list = []
+        prior_list = []
+        init_forward_h = torch.zeros((N, self.num_channels, H, W), device=seqn.device)
+        flatten_map = self.multi_prior_fusion(self.adists(seqn[:, 0]))
+        h_ = self.forward_rnn(torch.cat((seqn[:, 0], flatten_map, init_forward_h), dim=1))
+        hidden_list.append(h_)
+        prior_list.append(flatten_map)
+        #arb_up_kernel = self.kernel_predict(h_, coord, cell)
+        for j in range(1, self.count):
+            h_, _ = warp(h_, flow_queue[:, j-1])
+            flatten_map = self.multi_prior_fusion(self.adists(seqn[:, j]))
+            h_ = self.forward_rnn(torch.cat((seqn[:, j], flatten_map, h_), dim=1))
+            hidden_list.append(h_)
+            prior_list.append(flatten_map)
+
+        for i in range(self.count, T + self.count):
+            # if i == self.count:
+            #     for j in range(i, i+self.count):
+            #         h_, _ = warp(h_, flow_queue[:, j - 1])
+            #         h_ = self.forward_rnn(torch.cat((seqn[:, j], flatten_map[:, j], h_), dim=1))
+            #         hidden_list.append(h_)
+            # else:
+            #     #print(hidden_list[-1].shape)
+            #     #print(flow_queue[:, i + self.count - 1].shape)
+            h_, _ = warp(hidden_list[-1], flow_queue[:, i-1])
+            flatten_map = self.multi_prior_fusion(self.adists(seqn[:, i]))
+            h_ = self.forward_rnn(torch.cat((seqn[:, i], flatten_map, h_), dim=1))
+            hidden_list.append(h_)
+            prior_list.append(flatten_map)
+            if i > self.count:
+                hidden_list.pop(0)
+                prior_list.pop(0)
+
+            assert len(hidden_list) == self.count+1
+
+            refsr_flow_list = []
+            for j in range(1, self.count+1):
+                refsr_flow_list.append(extract_flow_torch(self.pwcnet, seqn[:, i-self.count], seqn[:, i-self.count + j]))
+
+            assert len(refsr_flow_list) == self.count
+
+            hidden = hidden_list[0]
+            fusion_h = self.predict(hidden, refsr_flow_list, hidden_list[1:], hidden.size()[-2:])
+
+            h = self.forward_rnn2(torch.cat((seqn[:, i-self.count], prior_list[0], fusion_h), dim=1))
+            res = self.d(torch.cat([h, hidden], dim=1))
+            seqdn[:, i-self.count] = self.upsample(res, arb_up_kernel, coord, seqn[:, i-self.count])
+            print(f"{seqdn.shape=}")
+
+        return seqdn

asvr_module/warp.py

@@ -0,0 +1,33 @@
+import torch
+
+def warp(x, flo):
+    '''
+    warp an image/tensor (im2) back to im1, according to the optical flow
+    x: [B, C, H, W] (im2)
+    flo: [B, 2, H, W] (flow)
+    '''
+    B, C, H, W = x.size()
+    # mesh grid
+    xx = torch.arange(0, W, device=x.device).view(1, -1).repeat(H, 1)
+    yy = torch.arange(0, H, device=x.device).view(-1, 1).repeat(1, W)
+    xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1)
+    yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1)
+    grid = torch.cat((xx, yy), 1).float()
+
+    if x.is_cuda:
+        grid = grid.to(x.device)
+    vgrid = torch.autograd.Variable(grid) + flo
+
+    # scale grid to [-1, 1]
+    vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :].clone() / max(W - 1, 1) - 1.0
+    vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :].clone() / max(H - 1, 1) - 1.0
+
+    vgrid = vgrid.permute(0, 2, 3, 1)
+    output = torch.nn.functional.grid_sample(x, vgrid, align_corners=True)
+    mask = torch.autograd.Variable(torch.ones((B, C, H, W), device=x.device))
+    mask = torch.nn.functional.grid_sample(mask, vgrid, align_corners=True)
+
+    mask[mask < 0.9999] = 0
+    mask[mask > 0] = 1
+
+    return output * mask, mask

asvr_run.py

@@ -0,0 +1,47 @@
+
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+import tqdm
+from comfy.utils import ProgressBar
+
+from .constants import ASVR_MODEL_PATH
+from .asvr_module.asvr_model import ASVR_Config, ASVR
+
+base_dir = Path(__file__).resolve().parent
+
+class ASVR_VideoSR:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_video": ("IMAGE",),
+                "height_scale": (
+                    "FLOAT",
+                    {"default": 8.0, "min": 0.1, "step": 0.01},
+                ),
+                "width_scale": (
+                    "FLOAT",
+                    {"default": 8.0, "min": 0.1, "step": 0.01},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("res_video",)
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes"
+
+    def todo(
+        self,
+        src_video: torch.Tensor,
+        height_scale: float,
+        width_scale: float
+    ):
+        asvr_runner = ASVR(ASVR_Config(
+            f"{base_dir}/{ASVR_MODEL_PATH}",
+            space_scale=(height_scale, width_scale)
+        ))
+        output = asvr_runner.inference(src_video)
+        return (output, )

constants.py

@@ -1,4 +1,5 @@
 WFEN_MODEL_PATH = "models/wfen/WFEN.pth"
 REALVIFORMER_MODEL_PATH = "models/realviformer/weights.pth"
 PROPIH_VGG_MODEL_PATH = "models/propih/vgg_normalised.pth"
-PROPIH_G_MODEL_PATH = "models/propih/latest_net_G.pth"
+PROPIH_G_MODEL_PATH = "models/propih/latest_net_G.pth"
+ASVR_MODEL_PATH = "models/asvr/refsrrnn_cuf_siren_adists_allstage_only_future_t2.pth"