Add Realviformer (#3)

README.md

@@ -2,7 +2,7 @@
 
 A collection of ComfyUI nodes. Including:
 1. WFEN Face Restore, [Efficient Face Super-Resolution via Wavelet-based Feature Enhancement Network](https://github.com/PRIS-CV/WFEN)
-
+2. [RealViformer](https://github.com/Yuehan717/RealViformer) - Investigating Attention for Real-World Video Super-Resolution
 
 ## Workflows
 
@@ -12,4 +12,19 @@ Download the model [here](https://github.com/PRIS-CV/WFEN?tab=readme-ov-file#get
 
 [workflow_wfen_facecrop.json](workflow/workflow_wfen_facecrop.json)
 
-![wfen_facecrop](workflow/wfen_facecrop.png)
+![wfen_facecrop](workflow/wfen_facecrop.png)
+
+### RealViformer
+
+Download the model [here](https://github.com/Yuehan717/RealViformer?tab=readme-ov-file#usage) and place it in `models/realviformer/weights.pth`.
+
+[workflow_realviformer.json](workflow/workflow_realviformer.json)
+
+![realviformer_example](example_realviformer.png)
+
+(Not a workflow-embedded image)
+
+https://github.com/user-attachments/assets/e89003c0-7be5-4263-b281-fd609807cea1
+
+RealViFormer upscale example
+

__init__.py

@@ -1,13 +1,14 @@
 from .run import (
-    WFEN,
+    WFEN, RealViFormerSR,
 )
 
 NODE_CLASS_MAPPINGS = {
-    "WFEN": WFEN
+    "WFEN": WFEN,
+    "RealViFormerSR": RealViFormerSR,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "WFEN": "WFEN Face Restore"
+    "RealViFormerSR": "RealViFormer Video SR"
 }
 
 

constants.py

@@ -1 +1,2 @@
-WFEN_MODEL_PATH = "models/wfen/WFEN.pth"
+WFEN_MODEL_PATH = "models/wfen/WFEN.pth"
+REALVIFORMER_MODEL_PATH = "models/realviformer/weights.pth"

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "fm_nodes"
-description = "A collection of ComfyUI nodes. Including: WFEN Face Restore, Efficient Face Super-Resolution via Wavelet-based Feature Enhancement Network"
-version = "1.0.0"
+description = "A collection of ComfyUI nodes. Including: WFEN Face Restore, RealViFormer Video Super Resolution"
+version = "1.0.1"
 license = {file = "LICENSE"}
 dependencies = ["torch", "einops"]
 

realviformer_module/arch_util.py

@@ -0,0 +1,583 @@
+import collections.abc
+import math
+import torch
+import warnings
+from itertools import repeat
+from torch import nn as nn
+from torch.nn import functional as F
+from torch.nn import init as init
+from torch.nn.modules.batchnorm import _BatchNorm
+
+from einops import rearrange
+import numbers
+
+
+@torch.no_grad() # @ is a decorator, which takes the decorated function as the argument of the function after @
+def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):
+    """Initialize network weights.
+
+    Args:
+        module_list (list[nn.Module] | nn.Module): Modules to be initialized.
+        scale (float): Scale initialized weights, especially for residual
+            blocks. Default: 1.
+        bias_fill (float): The value to fill bias. Default: 0
+        kwargs (dict): Other arguments for initialization function.
+    """
+    if not isinstance(module_list, list):
+        module_list = [module_list]
+    for module in module_list:
+        for m in module.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, nn.Linear):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, _BatchNorm):
+                init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+
+
+def make_layer(basic_block, num_basic_block, **kwarg):
+    """Make layers by stacking the same blocks.
+
+    Args:
+        basic_block (nn.module): nn.module class for basic block.
+        num_basic_block (int): number of blocks.
+
+    Returns:
+        nn.Sequential: Stacked blocks in nn.Sequential.
+    """
+    layers = []
+    for _ in range(num_basic_block):
+        layers.append(basic_block(**kwarg))
+    return nn.Sequential(*layers)
+
+
+class ResidualBlockNoBN(nn.Module):
+    """Residual block without BN.
+
+    It has a style of:
+        ---Conv-ReLU-Conv-+-
+         |________________|
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+            Default: 64.
+        res_scale (float): Residual scale. Default: 1.
+        pytorch_init (bool): If set to True, use pytorch default init,
+            otherwise, use default_init_weights. Default: False.
+    """
+
+    def __init__(self, num_feat=64, res_scale=1, pytorch_init=False):
+        super(ResidualBlockNoBN, self).__init__()
+        self.res_scale = res_scale
+        self.conv1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
+        self.conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
+        self.relu = nn.ReLU(inplace=True)
+
+        if not pytorch_init:
+            default_init_weights([self.conv1, self.conv2], 0.1)
+
+    def forward(self, x):
+        identity = x
+        out = self.conv2(self.relu(self.conv1(x)))
+        return identity + out * self.res_scale
+
+
+class Upsample(nn.Sequential):
+    """Upsample module.
+
+    Args:
+        scale (int): Scale factor. Supported scales: 2^n and 3.
+        num_feat (int): Channel number of intermediate features.
+    """
+
+    def __init__(self, scale, num_feat):
+        m = []
+        if (scale & (scale - 1)) == 0:  # scale = 2^n
+            for _ in range(int(math.log(scale, 2))):
+                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
+                m.append(nn.PixelShuffle(2))
+        elif scale == 3:
+            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
+            m.append(nn.PixelShuffle(3))
+        else:
+            raise ValueError(f'scale {scale} is not supported. Supported scales: 2^n and 3.')
+        super(Upsample, self).__init__(*m)
+
+
+def flow_warp(x, flow, interp_mode='bilinear', padding_mode='zeros', align_corners=True):
+    """Warp an image or feature map with optical flow.
+
+    Args:
+        x (Tensor): Tensor with size (n, c, h, w).
+        flow (Tensor): Tensor with size (n, h, w, 2), normal value.
+        interp_mode (str): 'nearest' or 'bilinear'. Default: 'bilinear'.
+        padding_mode (str): 'zeros' or 'border' or 'reflection'.
+            Default: 'zeros'.
+        align_corners (bool): Before pytorch 1.3, the default value is
+            align_corners=True. After pytorch 1.3, the default value is
+            align_corners=False. Here, we use the True as default.
+
+    Returns:
+        Tensor: Warped image or feature map.
+    """
+    assert x.size()[-2:] == flow.size()[1:3]
+    _, _, h, w = x.size()
+    # create mesh grid
+    grid_y, grid_x = torch.meshgrid(torch.arange(0, h).type_as(x), torch.arange(0, w).type_as(x))
+    grid = torch.stack((grid_x, grid_y), 2).float()  # W(x), H(y), 2 current position 
+    grid.requires_grad = False
+
+    vgrid = grid + flow # for all n batches do the same  current position --> warped position
+    # scale grid to [-1,1]
+    vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(w - 1, 1) - 1.0
+    vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(h - 1, 1) - 1.0
+    vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)
+    output = F.grid_sample(x, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode, align_corners=align_corners)
+
+    # TODO, what if align_corners=False
+    return output
+
+
+def resize_flow(flow, size_type, sizes, interp_mode='bilinear', align_corners=False):
+    """Resize a flow according to ratio or shape.
+
+    Args:
+        flow (Tensor): Precomputed flow. shape [N, 2, H, W].
+        size_type (str): 'ratio' or 'shape'.
+        sizes (list[int | float]): the ratio for resizing or the final output
+            shape.
+            1) The order of ratio should be [ratio_h, ratio_w]. For
+            downsampling, the ratio should be smaller than 1.0 (i.e., ratio
+            < 1.0). For upsampling, the ratio should be larger than 1.0 (i.e.,
+            ratio > 1.0).
+            2) The order of output_size should be [out_h, out_w].
+        interp_mode (str): The mode of interpolation for resizing.
+            Default: 'bilinear'.
+        align_corners (bool): Whether align corners. Default: False.
+
+    Returns:
+        Tensor: Resized flow.
+    """
+    _, _, flow_h, flow_w = flow.size()
+    if size_type == 'ratio':
+        output_h, output_w = int(flow_h * sizes[0]), int(flow_w * sizes[1])
+    elif size_type == 'shape':
+        output_h, output_w = sizes[0], sizes[1]
+    else:
+        raise ValueError(f'Size type should be ratio or shape, but got type {size_type}.')
+
+    input_flow = flow.clone()
+    ratio_h = output_h / flow_h
+    ratio_w = output_w / flow_w
+    input_flow[:, 0, :, :] *= ratio_w
+    input_flow[:, 1, :, :] *= ratio_h
+    resized_flow = F.interpolate(
+        input=input_flow, size=(output_h, output_w), mode=interp_mode) #, align_corners=align_corners)
+    return resized_flow
+
+
+# TODO: may write a cpp file
+def pixel_unshuffle(x, scale):
+    """ Pixel unshuffle.
+
+    Args:
+        x (Tensor): Input feature with shape (b, c, hh, hw).
+        scale (int): Downsample ratio.
+
+    Returns:
+        Tensor: the pixel unshuffled feature.
+    """
+    b, c, hh, hw = x.size()
+    out_channel = c * (scale**2)
+    assert hh % scale == 0 and hw % scale == 0
+    h = hh // scale
+    w = hw // scale
+    x_view = x.view(b, c, h, scale, w, scale)
+    return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)
+
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/weight_init.py
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            'mean is more than 2 std from [a, b] in nn.init.trunc_normal_. '
+            'The distribution of values may be incorrect.',
+            stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        low = norm_cdf((a - mean) / std)
+        up = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [low, up], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * low - 1, 2 * up - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution.
+
+    From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/weight_init.py
+
+    The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+# From PyTorch
+def _ntuple(n):
+
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
+
+class ConvResidualBlocks(nn.Module):
+    """Conv and residual block used in BasicVSR.
+
+    Args:
+        num_in_ch (int): Number of input channels. Default: 3.
+        num_out_ch (int): Number of output channels. Default: 64.
+        num_block (int): Number of residual blocks. Default: 15.
+    """
+
+    def __init__(self, num_in_ch=3, num_out_ch=64, num_block=15):
+        super().__init__()
+        self.main = nn.Sequential(
+            nn.Conv2d(num_in_ch, num_out_ch, 3, 1, 1, bias=True), nn.LeakyReLU(negative_slope=0.1, inplace=True),
+            make_layer(ResidualBlockNoBN, num_block, num_feat=num_out_ch))
+
+    def forward(self, fea):
+        return self.main(fea)
+    
+class ResidualBlocks(nn.Module):
+    """Residual block without shallow conv
+
+    Args:
+        num_in_ch (int): Number of input channels. Default: 3.
+        num_out_ch (int): Number of output channels. Default: 64.
+        num_block (int): Number of residual blocks. Default: 15.
+    """
+
+    def __init__(self, num_out_ch=64, num_block=15):
+        super().__init__()
+        self.main = make_layer(ResidualBlockNoBN, num_block, num_feat=num_out_ch)
+
+    def forward(self, fea):
+        return self.main(fea)
+
+class Conv2dwithActication(nn.Module):
+    """ Conv2d layers with activation
+    """
+    def __init__(self, num_in_ch, num_out_ch, num_layer=1):
+        super().__init__()
+        main = [nn.Conv2d(num_in_ch, num_out_ch, 3, 1, 1, bias=True), nn.LeakyReLU(negative_slope=0.1, inplace=True)]
+        for _ in range(num_layer-1):
+            main.append(nn.Conv2d(num_out_ch, num_out_ch, 3, 1, 1, bias=True))
+            main.append(nn.LeakyReLU(negative_slope=0.1, inplace=True))
+        self.main = nn.Sequential(*main)
+        
+    def forward(self, fea):
+        return self.main(fea)
+
+##########################################################################
+## LayerNorm
+
+def to_3d(x):
+    return rearrange(x, 'b c h w -> b (h w) c')
+
+def to_4d(x,h,w):
+    return rearrange(x, 'b (h w) c -> b c h w',h=h,w=w)
+
+class BiasFree_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(BiasFree_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return x / torch.sqrt(sigma+1e-5) * self.weight
+
+class WithBias_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(WithBias_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        mu = x.mean(-1, keepdim=True)
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return (x - mu) / torch.sqrt(sigma+1e-5) * self.weight + self.bias
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim, LayerNorm_type):
+        super(LayerNorm, self).__init__()
+        if LayerNorm_type =='BiasFree':
+            self.body = BiasFree_LayerNorm(dim)
+        else:
+            self.body = WithBias_LayerNorm(dim)
+
+    def forward(self, x):
+        h, w = x.shape[-2:]
+        return to_4d(self.body(to_3d(x)), h, w)
+
+
+##########################################################################
+## Multi-DConv Head Transposed Self-Attention (MDTA)
+class Attention(nn.Module):
+    def __init__(self, dim_q, dim_k, num_heads, bias):
+        super(Attention, self).__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.kv = nn.Conv2d(dim_k, dim_k*2, kernel_size=1, bias=bias)
+        self.kv_dwconv = nn.Conv2d(dim_k*2, dim_k*2, kernel_size=3, stride=1, padding=1, groups=dim_k*2, bias=bias)
+        self.q = nn.Conv2d(dim_q, dim_q, kernel_size=3, stride=1, padding=1, bias=bias)
+        self.project_out = nn.Conv2d(dim_q, dim_q, kernel_size=1, bias=bias)
+        
+
+
+    def forward(self, x, y):
+        b,c_kv,h,w = x.shape
+
+        kv = self.kv_dwconv(self.kv(x))
+        k,v = kv.chunk(2, dim=1)
+        
+        # c_q = y.size(1)
+        q = self.q(y)
+           
+        
+        q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature # c_q, c_kv (64, 32)
+        attn = attn.softmax(dim=-1)
+
+        out = (attn @ v)
+        
+        out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
+
+        out = self.project_out(out) # (b, 64, h, w)
+        return out
+
+
+##########################################################################
+# Attention block with layer norm
+class AttentionBlock(nn.Module):
+    def __init__(self, dim_q, dim_kv, num_heads, bias, LayerNorm_type=None, reduction='sum'):
+        super(AttentionBlock, self).__init__()
+        if LayerNorm_type is not None:
+            self.norm_q = LayerNorm(dim_q, LayerNorm_type)
+            self.norm_kv = LayerNorm(dim_kv, LayerNorm_type)
+        self.attn = Attention(dim_q, dim_kv, num_heads, bias)
+        if reduction == 'cat':
+            self.fusion = nn.Conv2d(dim_q*2, dim_q, kernel_size=1, bias=bias)
+    
+    def forward(self, x, y):
+        if hasattr(self, 'norm_kv'):
+            x = self.norm_kv(x)
+        if hasattr(self, 'norm_q'):
+            y = self.norm_q(y)
+        
+        v = self.attn(x, y) ## ()
+        
+        if hasattr(self, 'fusion'):
+            return self.fusion(torch.cat((y, v), dim=1))
+        else:
+            return y + v
+        
+
+##########################################################################
+## Simple attention function (parameter-free)
+    
+def cross_attention(x, ref):
+    b,c0,h,w = x.shape
+    _,c1,_,_ = ref.shape
+    q,k,v = x.clone(), ref.clone(), ref.clone()   
+    
+    q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=1)
+    k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=1)
+    v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=1)
+
+    q = torch.nn.functional.normalize(q, dim=-1)
+    k = torch.nn.functional.normalize(k, dim=-1)
+
+    attn = (q @ k.transpose(-2, -1)) * 1 # self.temperature [c0, c1]
+    max_, _ = torch.max(attn, dim=-1, keepdim=True)
+    attn_submax = attn - max_
+    attn = attn.softmax(dim=-1)
+    attn = attn * torch.lt(attn_submax, 0).int()
+    attn = attn.softmax(dim=-1)
+
+    out = (attn @ v) # [c0, hw]
+    
+    out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=1, h=h, w=w)
+
+    # out = self.project_out(out)
+    return out
+
+
+##########################################################################
+## Multi-DConv Head Cross-Attention (MDCA)
+class CrossChannelAttention(nn.Module):
+    def __init__(self, dim_q, dim_k, num_heads, bias):
+        super(CrossChannelAttention, self).__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.kv = nn.Conv2d(dim_k, dim_k*2, kernel_size=1, bias=bias)
+        self.kv_dwconv = nn.Conv2d(dim_k*2, dim_k*2, kernel_size=3, stride=1, padding=1, groups=dim_k*2, bias=bias)
+        self.q = nn.Conv2d(dim_q, dim_q, kernel_size=3, stride=1, padding=1, bias=bias)
+        self.project_out = nn.Conv2d(dim_q, dim_q, kernel_size=1, bias=bias)
+        
+    def forward(self, x, y, return_attn=False):
+        b,c_kv,h,w = x.shape
+
+        kv = self.kv_dwconv(self.kv(y))
+        k,v = kv.chunk(2, dim=1)
+        
+        # c_q = y.size(1)
+        q = self.q(x)
+           
+        q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature
+        
+        attn = attn.softmax(dim=-1)
+        out = (attn @ v)
+
+        out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
+
+        out = self.project_out(out) # (b, 64, h, w)
+        if return_attn:
+            return out, attn
+        else:
+            return out
+    
+##########################################################################
+# Channel-wise Attention block with layer norm
+class ChannelAttentionBlock(nn.Module):
+    def __init__(self, dim_q, dim_kv, dim_out,num_heads, bias, LayerNorm_type=None, 
+                 reduction=True, ch_compress=False, squeeze_factor=1):
+        super(ChannelAttentionBlock, self).__init__()
+        self.ch_compress = ch_compress and (squeeze_factor > 1)
+        if self.ch_compress:
+            if LayerNorm_type is not None:
+                self.norm_q = LayerNorm(dim_q//squeeze_factor, LayerNorm_type)
+                self.norm_kv = LayerNorm(dim_kv//squeeze_factor, LayerNorm_type)
+            self.attn = CrossChannelAttention(dim_q//squeeze_factor, dim_kv//squeeze_factor, num_heads, bias)
+        else:
+            if LayerNorm_type is not None:
+                self.norm_q = LayerNorm(dim_q, LayerNorm_type)
+                self.norm_kv = LayerNorm(dim_kv, LayerNorm_type)
+            self.attn = CrossChannelAttention(dim_q, dim_kv, num_heads, bias)
+        self.reduction = reduction
+        if reduction:
+            self.norm_out = LayerNorm(dim_kv, LayerNorm_type)
+            self.ffn = nn.Sequential(nn.Conv2d(dim_kv+dim_q, dim_out, kernel_size=1, bias=bias),
+                                    nn.Conv2d(dim_out, dim_out, kernel_size=3, stride=1, padding=1, groups=dim_out, bias=bias),
+                                    nn.GELU(),
+                                    nn.Conv2d(dim_out, dim_out, kernel_size=1, bias=bias))
+        
+        if self.ch_compress:
+            self.compress_q = nn.Conv2d(dim_q, dim_q//squeeze_factor, kernel_size=3, stride=1, padding=1, bias=bias)
+            self.compress_kv = nn.Conv2d(dim_kv, dim_kv//squeeze_factor, kernel_size=3, stride=1, padding=1, bias=bias)
+            self.expand = nn.Conv2d(dim_q//squeeze_factor, dim_q, kernel_size=3, stride=1, padding=1, bias=bias)
+    
+    def forward(self, x, y, return_attn=False):
+        
+        if self.ch_compress:
+            x_compressed = self.compress_q(x)
+            y_compressed = self.compress_kv(y)
+            if hasattr(self, 'norm_kv'):
+                y_compressed = self.norm_kv(y_compressed)
+            if hasattr(self, 'norm_q'):
+                x_compressed = self.norm_q(x_compressed)
+            v = self.expand(self.attn(x_compressed, y_compressed, return_attn))
+        else:
+            if hasattr(self, 'norm_kv'):
+                y = self.norm_kv(y)
+            if hasattr(self, 'norm_q'):
+                x = self.norm_q(x)
+            v = self.attn(x, y, return_attn) ## ()
+        if self.reduction:
+            out = self.ffn(torch.cat([x,self.norm_out(v)], dim=1))
+        return out

realviformer_module/realviformer_arch.py

@@ -0,0 +1,342 @@
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+from .arch_util import ResidualBlockNoBN, flow_warp, make_layer, LayerNorm, ChannelAttentionBlock
+from .spynet_arch import SpyNet
+from einops import rearrange
+import numbers
+##########################################################################
+## Layer Norm
+
+def to_3d(x):
+    return rearrange(x, 'b c h w -> b (h w) c')
+
+def to_4d(x,h,w):
+    return rearrange(x, 'b (h w) c -> b c h w',h=h,w=w)
+
+class BiasFree_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(BiasFree_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return x / torch.sqrt(sigma+1e-5) * self.weight
+
+class WithBias_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(WithBias_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        mu = x.mean(-1, keepdim=True)
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return (x - mu) / torch.sqrt(sigma+1e-5) * self.weight + self.bias
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim, LayerNorm_type):
+        super(LayerNorm, self).__init__()
+        if LayerNorm_type =='BiasFree':
+            self.body = BiasFree_LayerNorm(dim)
+        else:
+            self.body = WithBias_LayerNorm(dim)
+
+    def forward(self, x):
+        h, w = x.shape[-2:]
+        return to_4d(self.body(to_3d(x)), h, w)
+
+
+
+##########################################################################
+## Gated-Dconv Feed-Forward Network (GDFN)
+class FeedForward(nn.Module):
+    def __init__(self, dim, ffn_expansion_factor, bias):
+        super(FeedForward, self).__init__()
+
+        hidden_features = int(dim*ffn_expansion_factor)
+
+        self.project_in = nn.Conv2d(dim, hidden_features*2, kernel_size=1, bias=bias)
+
+        self.dwconv = nn.Conv2d(hidden_features*2, hidden_features*2, kernel_size=3, stride=1, padding=1, groups=hidden_features*2, bias=bias)
+
+        self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
+
+    def forward(self, x):
+        x = self.project_in(x)
+        x1, x2 = self.dwconv(x).chunk(2, dim=1)
+        x = F.gelu(x1) * x2
+        x = self.project_out(x)
+        return x
+
+
+
+##########################################################################
+## Multi-DConv Head Transposed Self-Attention (MDTA)
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads, bias, withmask=False):
+        super(Attention, self).__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.qkv = nn.Conv2d(dim, dim*3, kernel_size=1, bias=bias)
+        self.qkv_dwconv = nn.Conv2d(dim*3, dim*3, kernel_size=3, stride=1, padding=1, groups=dim*3, bias=bias)
+        self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
+        self.withmask = withmask
+        if withmask:
+            self.max_dsc = nn.AdaptiveMaxPool1d(1)
+            self.avg_dsc = nn.AdaptiveAvgPool1d(1)
+            self.linear1 = nn.Linear(2, 1, bias=bias)
+            self.linear2 = nn.Linear(dim//num_heads, dim//num_heads, bias=bias)
+
+    def forward(self, x):
+        b,c,h,w = x.shape
+
+        qkv = self.qkv_dwconv(self.qkv(x))
+        q,k,v = qkv.chunk(3, dim=1)   
+        
+        q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+    
+            
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature # b, h, c, c
+        
+        if self.withmask:
+            mask_in = attn.clone()
+            b_, h_, c_, _ = mask_in.shape 
+            dsc = torch.cat([self.max_dsc(mask_in.reshape(b_, h_*c_, c_)), self.avg_dsc(mask_in.reshape((b_, h_*c_, c_)))], dim=-1).reshape(b_, h_, c_, 2) # b,h,c,2
+            mask = self.linear1(dsc) # b,h,c,1
+            mask = self.linear2(F.gelu(mask).transpose(-2,-1)).transpose(-2,-1) # b, h, c, 1
+            mask = F.sigmoid(mask)
+        
+        attn = attn.softmax(dim=-1)
+        out = (attn @ v)
+        
+        if self.withmask:
+            out = out * mask
+        
+        out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
+
+        out = self.project_out(out)
+        return out
+
+
+
+##########################################################################
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, num_heads, ffn_expansion_factor, bias, LayerNorm_type, 
+                 masked=False, squeeze_factor=1.0, ch_compress=False):
+        super(TransformerBlock, self).__init__()
+        if ch_compress:
+            self.norm1 = LayerNorm(dim//squeeze_factor, LayerNorm_type)
+            self.attn = Attention(dim//squeeze_factor, num_heads, bias, masked)
+        else:
+            self.norm1 = LayerNorm(dim, LayerNorm_type)
+            self.attn = Attention(dim, num_heads, bias, masked)
+        self.norm2 = LayerNorm(dim, LayerNorm_type)
+        self.ffn = FeedForward(dim, ffn_expansion_factor, bias)
+        self.ch_compress = ch_compress and (squeeze_factor > 1)
+        if self.ch_compress:
+            self.compress = nn.Conv2d(dim, dim//squeeze_factor, kernel_size=3, stride=1, padding=1, bias=bias)
+            self.expand = nn.Conv2d(dim//squeeze_factor, dim, kernel_size=3, stride=1, padding=1, bias=bias)
+
+    def forward(self, x):
+        if self.ch_compress:
+            x_compressed = self.compress(x)
+            x = x +  self.expand(self.attn(self.norm1(x_compressed)))
+        else:
+            x = x + self.attn(self.norm1(x))
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+##########################################################################
+## Resizing modules
+class Downsample(nn.Module):
+    def __init__(self, n_feat):
+        super(Downsample, self).__init__()
+
+        self.body = nn.Sequential(nn.Conv2d(n_feat, n_feat//2, kernel_size=3, stride=1, padding=1, bias=False),
+                                  nn.PixelUnshuffle(2))
+
+    def forward(self, x):
+        return self.body(x)
+
+class Upsample(nn.Module):
+    def __init__(self, n_feat):
+        super(Upsample, self).__init__()
+
+        self.body = nn.Sequential(nn.Conv2d(n_feat, n_feat*2, kernel_size=3, stride=1, padding=1, bias=False),
+                                  nn.PixelShuffle(2))
+
+    def forward(self, x):
+        return self.body(x)
+
+class RealViformer(nn.Module):
+    """A recurrent network for video SR. Now only x4 is supported.
+
+    Args:
+        num_feat (int): Number of channels. Default: 64.
+        num_block (int): Number of residual blocks for each branch. Default: 15
+        spynet_path (str): Path to the pretrained weights of SPyNet. Default: None.
+    """
+
+    def __init__(self, num_feat=64, num_blocks=[], spynet_path=None, heads=[], ffn_expansion_factor=2, masked=False,
+                 merge_head=1, merge_compress=False, merge_compress_factor=1.0,
+                 bias=False, LayerNorm_type='BiasFree', ch_compress=False, squeeze_factor=[1,1,1]):
+        super().__init__()
+        # embedding
+        self.num_feat = num_feat
+        self.shallow_extraction = nn.Sequential(nn.Conv2d(3, num_feat, 3, 1, 1),
+                                                TransformerBlock(dim=num_feat, num_heads=heads[0], ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type, masked=False, ch_compress=False, squeeze_factor=squeeze_factor[0]))
+        
+        # alignment
+        self.spynet = SpyNet(spynet_path)
+        self.attn_merge = ChannelAttentionBlock(dim_q= num_feat, dim_kv = num_feat, dim_out=num_feat, num_heads=merge_head, bias=False,
+                                                LayerNorm_type='BiasFree', reduction=True, 
+                                                ch_compress=merge_compress, squeeze_factor=merge_compress_factor)
+        
+        # propagation
+        # self.forward_trunk = ConvResidualBlocks(num_feat * 2, num_feat, num_block-2)
+        ## dim = 48, num_heads = 1, num_block = 4
+        self.encoder_level1 = nn.Sequential(*[TransformerBlock(dim=num_feat, num_heads=heads[0], ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type, masked=masked,ch_compress=ch_compress, squeeze_factor=squeeze_factor[0]) for i in range(num_blocks[0])])
+        ## dim=96, num_heads = 2, num_block=6
+        self.down1_2 = Downsample(num_feat) ## From Level 1 to Level 2
+        self.encoder_level2 = nn.Sequential(*[TransformerBlock(dim=int(num_feat*2**1), num_heads=heads[1], ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type, masked=masked,ch_compress=ch_compress, squeeze_factor=squeeze_factor[0]) for i in range(num_blocks[1])])
+        ## dim=192, num_heads = 4, num_block=6
+        self.down2_3 = Downsample(int(num_feat*2**1)) ## From Level 2 to Level 3
+        self.latent= nn.Sequential(*[TransformerBlock(dim=int(num_feat*2**2), num_heads=heads[2], ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type, masked=masked,ch_compress=ch_compress, squeeze_factor=squeeze_factor[0]) for i in range(num_blocks[2])])
+        ## dim=96, num_heads = 2, num_block=6
+        self.up3_2 = Upsample(int(num_feat*2**2)) ## From Level 3 to Level 2
+        self.reduce_chan_level2 = nn.Conv2d(int(num_feat*2**2), int(num_feat*2**1), kernel_size=1, bias=bias)
+        self.decoder_level2 = nn.Sequential(*[TransformerBlock(dim=int(num_feat*2**1), num_heads=heads[1], ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type, masked=masked,ch_compress=ch_compress, squeeze_factor=squeeze_factor[0]) for i in range(num_blocks[1])])
+        ## dim=96, num_heads = 1, num_block=4
+        self.up2_1 = Upsample(int(num_feat*2**1))  ## From Level 2 to Level 1  (NO 1x1 conv to reduce channels)
+        self.decoder_level1 = nn.Sequential(*[TransformerBlock(dim=int(num_feat*2**1), num_heads=heads[0], ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type, masked=masked,ch_compress=ch_compress, squeeze_factor=squeeze_factor[0]) for i in range(num_blocks[0])])
+        ## dim=96, num_heads = 1, num_block=2
+        self.refinement = nn.Sequential(*[TransformerBlock(dim=int(num_feat*2**1), num_heads=heads[0], ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type, masked=masked,ch_compress=ch_compress, squeeze_factor=squeeze_factor[0]) for i in range(num_blocks[3])])
+        
+        # reconstruction
+        self.compress = nn.Conv2d(num_feat*2**1, num_feat, kernel_size=1, bias=True)
+        # self.fusion = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
+        self.upconv1 = nn.Conv2d(num_feat*2**1, num_feat * 4, 3, 1, 1, bias=True)
+        self.upconv2 = nn.Conv2d(num_feat, 64 * 4, 3, 1, 1, bias=True)
+        self.conv_hr = nn.Conv2d(64, 64, 3, 1, 1)
+        self.conv_last = nn.Conv2d(64, 3, 3, 1, 1)
+
+        self.pixel_shuffle = nn.PixelShuffle(2)
+
+        # activation functions
+        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+
+
+    def get_flow(self, x):
+        
+        b, n, c, h, w = x.size()
+
+        x_1 = x[:, :-1, :, :, :].reshape(-1, c, h, w)
+        x_2 = x[:, 1:, :, :, :].reshape(-1, c, h, w)
+        
+        # print(x_1.size(), x_2.size())
+
+        flows_forward = self.spynet(x_2, x_1).view(b, n - 1, 2, h, w)
+        # flows_forward = self.spynet(x_2, x_1).view(b, n - 1, 2, h, w)
+        
+        return flows_forward
+     
+    # @torch.compile
+    def forward(self, x, current_iter=None):
+        """Forward function of BasicVSR.
+
+        Args:
+            x: Input frames with shape (b, n, c, h, w). n is the temporal dimension / number of frames.
+        """
+        b, n, _, h, w = x.size()
+        # self.check_if_mirror_extended(x)
+        flows_forward = self.get_flow(x.clone())
+        
+        out_l = []
+        # forward branch
+        # feat_prop = torch.zeros_like(feat_prop)
+        feat_prop = x.new_zeros(b, self.num_feat, h, w)
+        for i in range(0, n):
+            x_i = x[:, i, :, :, :]
+            feat_shallow = self.shallow_extraction(x_i)
+            if i > 0:
+                flow = flows_forward[:, i - 1, :, :, :]
+                feat_prop = flow_warp(feat_prop, flow.permute(0, 2, 3, 1))
+            ## alignment ##
+            feat_prop = self.attn_merge(feat_shallow, feat_prop)
+            # if hasattr(self, 'align_correction'):
+            #     feat_prop = self.align_correction(feat_shallow, feat_prop)            
+            
+            ## Propagation ##
+            out_enc_level1 = self.encoder_level1(feat_prop)
+        
+            inp_enc_level2 = self.down1_2(out_enc_level1)
+            out_enc_level2 = self.encoder_level2(inp_enc_level2)
+
+            inp_enc_level3 = self.down2_3(out_enc_level2)     
+            latent = self.latent(inp_enc_level3) 
+
+            inp_dec_level2 = self.up3_2(latent)
+            inp_dec_level2 = torch.cat([inp_dec_level2, out_enc_level2], 1)
+            inp_dec_level2 = self.reduce_chan_level2(inp_dec_level2)
+            out_dec_level2 = self.decoder_level2(inp_dec_level2) 
+
+            inp_dec_level1 = self.up2_1(out_dec_level2)
+            inp_dec_level1 = torch.cat([inp_dec_level1, out_enc_level1], 1)
+            out_dec_level1 = self.decoder_level1(inp_dec_level1)
+            
+            out = self.refinement(out_dec_level1)
+            feat_prop = self.compress(out)
+            # upsample
+            out = self.lrelu(self.pixel_shuffle(self.upconv1(out)))
+            out = self.lrelu(self.pixel_shuffle(self.upconv2(out)))
+            out = self.lrelu(self.conv_hr(out))
+            # base = F.interpolate(x_i, scale_factor=4, mode='bilinear', align_corners=False)
+            out = self.conv_last(out)
+            base = F.interpolate(x_i, scale_factor=4, mode='bilinear', align_corners=False)
+            out += base
+            
+            out_l.append(out)
+
+        return torch.stack(out_l, dim=1)
+
+class ConvResidualBlocks(nn.Module):
+    """Conv and residual block used in BasicVSR.
+
+    Args:
+        num_in_ch (int): Number of input channels. Default: 3.
+        num_out_ch (int): Number of output channels. Default: 64.
+        num_block (int): Number of residual blocks. Default: 15.
+    """
+
+    def __init__(self, num_in_ch=3, num_out_ch=64, num_block=15):
+        super().__init__()
+        self.main = nn.Sequential(
+            nn.Conv2d(num_in_ch, num_out_ch, 3, 1, 1, bias=True), nn.LeakyReLU(negative_slope=0.1, inplace=True),
+            make_layer(ResidualBlockNoBN, num_block, num_feat=num_out_ch))
+
+    def forward(self, fea):
+        return self.main(fea)
+

realviformer_module/spynet_arch.py

@@ -0,0 +1,114 @@
+import math
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+from .arch_util import flow_warp
+
+
+class BasicModule(nn.Module):
+    """Basic Module for SpyNet.
+    """
+
+    def __init__(self):
+        super(BasicModule, self).__init__()
+
+        self.basic_module = nn.Sequential(
+            nn.Conv2d(in_channels=8, out_channels=32, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False),
+            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False),
+            nn.Conv2d(in_channels=64, out_channels=32, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False),
+            nn.Conv2d(in_channels=32, out_channels=16, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False),
+            nn.Conv2d(in_channels=16, out_channels=2, kernel_size=7, stride=1, padding=3))
+
+    def forward(self, tensor_input):
+        return self.basic_module(tensor_input)
+
+class SpyNet(nn.Module):
+    """SpyNet architecture.
+
+    Args:
+        load_path (str): path for pretrained SpyNet. Default: None.
+    """
+
+    def __init__(self, load_path=None):
+        super(SpyNet, self).__init__()
+        self.basic_module = nn.ModuleList([BasicModule() for _ in range(6)])
+        if load_path:
+            self.load_state_dict(torch.load(load_path, map_location=lambda storage, loc: storage)['params'])
+
+        self.register_buffer('mean', torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+        self.register_buffer('std', torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+
+    def preprocess(self, tensor_input):
+        tensor_output = (tensor_input - self.mean) / self.std
+        return tensor_output
+
+    def process(self, ref, supp):
+        flow = []
+        W = ref.shape[-1]
+
+        ref = [self.preprocess(ref)]
+        supp = [self.preprocess(supp)]
+        
+        
+        r = 5 if W > 2**5 else int(math.log2(W) -1)
+        # print(W,r)
+
+        for level in range(r):
+            # print(ref[0].size())
+            ref.insert(0, F.avg_pool2d(input=ref[0], kernel_size=2, stride=2, count_include_pad=False))
+            supp.insert(0, F.avg_pool2d(input=supp[0], kernel_size=2, stride=2, count_include_pad=False))
+
+        flow = ref[0].new_zeros(
+            [ref[0].size(0), 2,
+             int(math.floor(ref[0].size(2) / 2.0)),
+             int(math.floor(ref[0].size(3) / 2.0))])
+        
+        
+
+        for level in range(len(ref)):
+            upsampled_flow = F.interpolate(input=flow, scale_factor=2, mode='bilinear', align_corners=True) * 2.0
+
+            if upsampled_flow.size(2) != ref[level].size(2):
+                upsampled_flow = F.pad(input=upsampled_flow, pad=[0, 0, 0, 1], mode='replicate')
+            if upsampled_flow.size(3) != ref[level].size(3):
+                upsampled_flow = F.pad(input=upsampled_flow, pad=[0, 1, 0, 0], mode='replicate')
+
+            flow = self.basic_module[level](torch.cat([
+                ref[level],
+                flow_warp(
+                    supp[level], upsampled_flow.permute(0, 2, 3, 1), interp_mode='bilinear', padding_mode='border'),
+                upsampled_flow
+            ], 1)) + upsampled_flow
+
+        return flow
+
+    def forward(self, ref, supp):
+        assert ref.size() == supp.size()
+
+        h, w = ref.size(2), ref.size(3)
+        w_floor = math.floor(math.ceil(w / 32.0) * 32.0)
+        h_floor = math.floor(math.ceil(h / 32.0) * 32.0)
+
+        ref = F.interpolate(input=ref, size=(h_floor, w_floor), mode='bilinear', align_corners=False)
+        supp = F.interpolate(input=supp, size=(h_floor, w_floor), mode='bilinear', align_corners=False)
+
+        flow = F.interpolate(input=self.process(ref, supp), size=(h, w), mode='bilinear', align_corners=False)
+
+        flow[:, 0, :, :] *= float(w) / float(w_floor)
+        flow[:, 1, :, :] *= float(h) / float(h_floor)
+
+        return flow
+    
+    def get_flow(self, x):
+        b, n, c, h, w = x.size()
+
+        x_1 = x[:, :-1, :, :, :].reshape(-1, c, h, w)
+        x_2 = x[:, 1:, :, :, :].reshape(-1, c, h, w)
+
+        flows_backward = self.forward(x_1, x_2).view(b, n - 1, 2, h, w)
+
+        flows_forward = self.forward(x_2, x_1).view(b, n - 1, 2, h, w)
+
+        return flows_forward, flows_backward
+

run.py

@@ -6,10 +6,14 @@
 """
 
 from pathlib import Path
+import tqdm
 
 import torch
+import torch.nn.functional as F
+from comfy.utils import ProgressBar
 
-from .constants import WFEN_MODEL_PATH
+from .constants import REALVIFORMER_MODEL_PATH, WFEN_MODEL_PATH
+from .realviformer_module.realviformer_arch import RealViformer
 from .utils import ensure_size
 from .wfen_module.wfen_model import WFENModel
 
@@ -41,8 +45,91 @@ class WFEN:
         wfening.for_load_pretrain_model(f"{base_dir}/{WFEN_MODEL_PATH}")
 
         out_tensor = wfening.netG.forward(source_tensor)
-        # out_tensor = ensure_size(out_tensor, size=src_img.shape[1:3])
         print(f"{out_tensor.shape=}")
         out_tensor = out_tensor.permute(0, 2, 3, 1)
 
         return (out_tensor,)
+
+
+def inference_realviformer(images: torch.Tensor, model: RealViformer) -> torch.Tensor:
+    padded = False
+    with torch.no_grad():
+        h, w = images.shape[-2:]
+        ah, aw = h % 4, w % 4
+        padh = 0 if ah == 0 else 4 - ah
+        padw = 0 if aw == 0 else 4 - aw
+        if padh != 0 or padw != 0:
+            padded = True
+            images = F.pad(
+                images.squeeze(0), pad=(padw, 0, padh, 0), mode="reflect"
+            ).unsqueeze(0)
+        outputs = model(images)
+
+    if padded:
+        outputs = outputs[..., padh * 4 :, padw * 4 :]
+
+    return outputs
+
+
+class RealViFormerSR:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_video": ("IMAGE",),
+                "interval": (
+                    "INT",
+                    {"default": 50, "min": 0, "step": 1},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("res_video",)
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes"
+
+    def todo(self, src_video: torch.Tensor, interval: int):
+        src_video = src_video.permute(0, 3, 1, 2)
+        print(f"{src_video.shape=}")
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = RealViformer(
+            num_feat=48,
+            num_blocks=[2, 3, 4, 1],
+            spynet_path=None,
+            heads=[1, 2, 4],
+            ffn_expansion_factor=2.66,
+            merge_head=2,
+            bias=False,
+            LayerNorm_type="BiasFree",
+            ch_compress=True,
+            squeeze_factor=[4, 4, 4],
+            masked=True,
+        )
+        model.load_state_dict(
+            torch.load(f"{base_dir}/{REALVIFORMER_MODEL_PATH}")["params"], strict=False
+        )
+        model.eval()
+        model = model.to(device)
+
+        if src_video.shape[0] <= interval:
+            out_tensor = inference_realviformer(
+                src_video.unsqueeze(0).to(device), model
+            )
+            out_tensor = out_tensor.squeeze(dim=0).permute(0, 2, 3, 1)
+            return (out_tensor,)
+
+        num_imgs = src_video.shape[0]
+        outputs: list[torch.Tensor] = []
+        pbar = ProgressBar(num_imgs)
+
+        for idx in tqdm.tqdm(range(0, num_imgs, interval)):
+            interval = min(interval, num_imgs - idx)
+            imgs = src_video[idx : idx + interval]
+            imgs = imgs.unsqueeze(0).to(device)  # [b, n, c, h, w]
+            outputs.append(inference_realviformer(imgs, model).squeeze(dim=0))
+            pbar.update_absolute(idx + interval, num_imgs)
+
+        out_tensor = torch.cat(outputs, dim=0).permute(0, 2, 3, 1)
+        return (out_tensor,)

workflow/workflow_realviformer.json

@@ -0,0 +1,327 @@
+{
+  "last_node_id": 5,
+  "last_link_id": 5,
+  "nodes": [
+    {
+      "id": 5,
+      "type": "VHS_VideoInfoLoaded",
+      "pos": [
+        746,
+        316
+      ],
+      "size": {
+        "0": 304.79998779296875,
+        "1": 106
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "link": 2
+        }
+      ],
+      "outputs": [
+        {
+          "name": "fps🟦",
+          "type": "FLOAT",
+          "links": [
+            3
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "frame_count🟦",
+          "type": "INT",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "duration🟦",
+          "type": "FLOAT",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "width🟦",
+          "type": "INT",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "height🟦",
+          "type": "INT",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoInfoLoaded"
+      },
+      "widgets_values": {}
+    },
+    {
+      "id": 1,
+      "type": "VHS_LoadVideo",
+      "pos": [
+        407,
+        204
+      ],
+      "size": [
+        235.1999969482422,
+        398.0031478761688
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            1
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [
+            5
+          ],
+          "shape": 3,
+          "slot_index": 2
+        },
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "links": [
+            2
+          ],
+          "shape": 3,
+          "slot_index": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_LoadVideo"
+      },
+      "widgets_values": {
+        "video": "cat_lq.mp4",
+        "force_rate": 0,
+        "force_size": "Disabled",
+        "custom_width": 512,
+        "custom_height": 512,
+        "frame_load_cap": 0,
+        "skip_first_frames": 0,
+        "select_every_nth": 1,
+        "choose video to upload": "image",
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "frame_load_cap": 0,
+            "skip_first_frames": 0,
+            "force_rate": 0,
+            "filename": "cat_lq.mp4",
+            "type": "input",
+            "format": "video/mp4",
+            "select_every_nth": 1
+          }
+        }
+      }
+    },
+    {
+      "id": 4,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        1139,
+        228
+      ],
+      "size": [
+        315,
+        487
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 4
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 5
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        },
+        {
+          "name": "frame_rate",
+          "type": "FLOAT",
+          "link": 3,
+          "widget": {
+            "name": "frame_rate"
+          }
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 8,
+        "loop_count": 0,
+        "filename_prefix": "VideoSR",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 19,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": true,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "VideoSR_00003.mp4",
+            "subfolder": "",
+            "type": "output",
+            "format": "video/h264-mp4",
+            "frame_rate": 20
+          }
+        }
+      }
+    },
+    {
+      "id": 2,
+      "type": "RealViFormerSR",
+      "pos": [
+        741,
+        173
+      ],
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "src_video",
+          "type": "IMAGE",
+          "link": 1
+        }
+      ],
+      "outputs": [
+        {
+          "name": "res_video",
+          "type": "IMAGE",
+          "links": [
+            4
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "RealViFormerSR"
+      },
+      "widgets_values": [
+        20
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      1,
+      0,
+      2,
+      0,
+      "IMAGE"
+    ],
+    [
+      2,
+      1,
+      3,
+      5,
+      0,
+      "VHS_VIDEOINFO"
+    ],
+    [
+      3,
+      5,
+      0,
+      4,
+      4,
+      "FLOAT"
+    ],
+    [
+      4,
+      2,
+      0,
+      4,
+      0,
+      "IMAGE"
+    ],
+    [
+      5,
+      1,
+      2,
+      4,
+      1,
+      "AUDIO"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1.1000000000000008,
+      "offset": [
+        -71.51282852004415,
+        26.7761769988706
+      ]
+    }
+  },
+  "version": 0.4
+}