Add vfi mamba (#6)

.gitignore

@@ -1,4 +1,5 @@
 *.pth
+*.pkl
 
 # Byte-compiled / optimized / DLL files
 __pycache__

README.md

@@ -7,7 +7,8 @@ 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. [**VFIMamba**](#vfimamba). Paper: [Video Frame Interpolation with State Space Models](https://github.com/MCG-NJU/VFIMamba)
+
 ## Workflows
 
 ### WFEN
@@ -52,3 +53,25 @@ 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)
+
+### VFIMamba
+
+Download the models from the [huggingface page](https://huggingface.co/MCG-NJU/VFIMamba/tree/main/ckpt)
+
+```
+models/vfimamba/VFIMamba_S.pkl
+models/vfimamba/VFIMamba.pkl
+```
+
+You will need to install `mamba-ssm`, which does not have a prebuilt Windows binary. You will need:
+1. [triton](https://github.com/triton-lang/triton). Prebuilt for `Python 3.10 and 3.11` can be found here: https://github.com/triton-lang/triton/issues/2881 - https://huggingface.co/madbuda/triton-windows-builds/tree/main
+2. [causal-conv1d](https://github.com/Dao-AILab/causal-conv1d). Follow this post: https://github.com/NVlabs/MambaVision/issues/14#issuecomment-2232581078
+3. [mamba-ssm](https://github.com/state-spaces/mamba). Follow this tutorial: https://blog.csdn.net/yyywxk/article/details/140420538. Fork that followed all the steps: https://github.com/FuouM/mamba-windows-build
+
+I've built `mamba-ssm` for `Python 3.11, torch 2.3.0+cu121`, which can be obtained here: https://huggingface.co/FuouM/mamba-ssm-windows-builds/tree/main
+
+To install, `pip install [].whl`
+
+[workflow_vfi_mamba.json](workflow/workflow_vfi_mamba.json)
+
+![example_vfi_mamba](example_vfi_mamba.png)

__init__.py

@@ -16,5 +16,13 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CoLIE_LowLight_Enhance": "CoLIE LowLight Enhance",
 }
 
+try:
+    from .vfi_mamba_run import VFI_MAMBA
+    NODE_CLASS_MAPPINGS["VFI_MAMBA"] = VFI_MAMBA
+    NODE_DISPLAY_NAME_MAPPINGS["VFI_MAMBA"] = "VFI Mamba"
+except ImportError as e:
+    print(f"Failed to load VFI_MAMBA. This will not affect other nodes. {e}")
+    pass
+
 
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

constants.py

@@ -1,4 +1,10 @@
 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"
+VFI_MAMBA_MODELS = ["VFIMamba", "VFIMamba_S"]
+VFI_MAMBA_PATHS = {
+    "VFIMamba": "models/vfimamba/VFIMamba.pkl",
+    "VFIMamba_S": "models/vfimamba/VFIMamba_S.pkl",
+}
+VFI_MAMBA_DEFAULT = "VFIMamba"

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "fm_nodes"
-description = "An expanding collection of ComfyUI nodes focused on new papers. Including: WFEN, RealViFormer, ProPIH, CoLIE"
-version = "1.0.3"
+description = "An expanding collection of ComfyUI nodes focused on new papers. Including: WFEN, RealViFormer, ProPIH, CoLIE, VFIMamba"
+version = "1.0.4"
 license = {file = "LICENSE"}
 dependencies = ["torch", "einops"]
 

vfi_mamba_run.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+
+import torch
+
+from .constants import VFI_MAMBA_DEFAULT, VFI_MAMBA_MODELS, VFI_MAMBA_PATHS
+from .vfimamba_module.vfi_mamba_model import run_vfi_mamba
+
+base_dir = Path(__file__).resolve().parent
+
+
+class VFI_MAMBA:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_video": ("IMAGE",),
+                "vfi_scale": (
+                    "INT",
+                    {"default": 2, "min": 2, "step": 1},
+                ),
+                "model": (VFI_MAMBA_MODELS, {"default": VFI_MAMBA_DEFAULT}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("res_video",)
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes"
+
+    def todo(
+        self, src_video: torch.Tensor, vfi_scale: int, model: str
+    ):
+        return (
+            run_vfi_mamba(
+                src_video=src_video,
+                model_name=model,
+                model_path=f"{base_dir}/{VFI_MAMBA_PATHS[model]}",
+                vfi_scale=vfi_scale,
+                scale=0,
+            ),
+        )

vfimamba_module/Trainer_finetune.py

@@ -0,0 +1,108 @@
+import torch
+import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from .config import LOCAL
+
+
+class Model:
+    def __init__(self, local_rank, MODEL_CONFIG):
+        backbonetype, multiscaletype = MODEL_CONFIG["MODEL_TYPE"]
+        backbonecfg, multiscalecfg = MODEL_CONFIG["MODEL_ARCH"]
+        self.net = multiscaletype(backbonetype(**backbonecfg), **multiscalecfg)
+        self.name = MODEL_CONFIG["LOGNAME"]
+        self.device()
+
+        self.local = LOCAL
+
+        # train
+        if local_rank != -1:
+            self.net = DDP(self.net, device_ids=[local_rank], output_device=local_rank)
+
+    def train(self):
+        self.net.train()
+
+    def eval(self):
+        self.net.eval()
+
+    def device(self):
+        self.net.to(torch.device("cuda"))
+
+    def load_model(self, model_path: str, name=None, rank=0, real=False):
+        def convert(param):
+            return {
+                k.replace("module.", ""): v
+                for k, v in param.items()
+                if "module." in k and "attn_mask" not in k and "HW" not in k
+            }
+
+        if rank <= 0:
+            if name is None:
+                name = self.name
+            print(f"loading {name} ckpt")
+            self.net.load_state_dict(convert(torch.load(f"{model_path}")), strict=True)
+
+    @torch.no_grad()
+    def hr_inference(
+        self, img0, img1, local, TTA=False, down_scale=1.0, timestep=0.5, fast_TTA=False
+    ):
+        """
+        Infer with down_scale flow
+        Noting: return BxCxHxW
+        """
+
+        def infer(imgs):
+            img0, img1 = imgs[:, :3], imgs[:, 3:6]
+            imgs_down = F.interpolate(
+                imgs, scale_factor=down_scale, mode="bilinear", align_corners=False
+            )
+
+            flow, mask = self.net.calculate_flow(imgs_down, timestep, local=local)
+
+            flow = F.interpolate(
+                flow, scale_factor=1 / down_scale, mode="bilinear", align_corners=False
+            ) * (1 / down_scale)
+            mask = F.interpolate(
+                mask, scale_factor=1 / down_scale, mode="bilinear", align_corners=False
+            )
+
+            af = self.net.feature_bone(img0, img1)
+            pred = self.net.coraseWarp_and_Refine(imgs, af, flow, mask)
+            return pred
+
+        imgs = torch.cat((img0, img1), 1)
+        if fast_TTA:
+            imgs_ = imgs.flip(2).flip(3)
+            input = torch.cat((imgs, imgs_), 0)
+            preds = infer(input)
+            return (preds[0] + preds[1].flip(1).flip(2)).unsqueeze(0) / 2.0
+
+        if not TTA:
+            return infer(imgs)
+        else:
+            return (infer(imgs) + infer(imgs.flip(2).flip(3)).flip(2).flip(3)) / 2
+
+    @torch.no_grad()
+    def inference(
+        self, img0, img1, local, TTA=False, timestep=0.5, scale=0, fast_TTA=False
+    ):
+        imgs = torch.cat((img0, img1), 1)
+        """
+        Noting: return BxCxHxW
+        """
+        if fast_TTA:
+            imgs_ = imgs.flip(2).flip(3)
+            input = torch.cat((imgs, imgs_), 0)
+            _, _, _, preds = self.net(
+                input, local=local, timestep=timestep, scale=scale
+            )
+            return (preds[0] + preds[1].flip(1).flip(2)).unsqueeze(0) / 2.0
+
+        _, _, _, pred = self.net(imgs, timestep=timestep, scale=scale, local=local)
+        if not TTA:
+            return pred
+        else:
+            _, _, _, pred2 = self.net(
+                imgs.flip(2).flip(3), timestep=timestep, scale=scale, local=local
+            )
+            return (pred + pred2.flip(2).flip(3)) / 2

vfimamba_module/config.py

@@ -0,0 +1,50 @@
+from functools import partial
+
+import torch.nn as nn
+
+from .feature_extractor import feature_extractor as mamba_extractor
+from .flow_estimation import MultiScaleFlow as mamba_estimation
+
+LOG = "VFIMamba_S"
+LOCAL = 2
+
+
+# '''==========Model config=========='''
+def init_model_config(F=32, W=7, depth=[2, 2, 2, 4, 4], M=False):
+    """This function should not be modified"""
+    return {
+        "embed_dims": [(2**i) * F for i in range(len(depth))],
+        "motion_dims": [0, 0, 0, 8 * F // depth[-2], 16 * F // depth[-1]],
+        "num_heads": [8 * (2**i) * F // 32 for i in range(len(depth) - 3)],
+        "mlp_ratios": [4 for i in range(len(depth) - 3)],
+        "qkv_bias": True,
+        "norm_layer": partial(nn.LayerNorm, eps=1e-6),
+        "depths": depth,
+        "window_sizes": [W for i in range(len(depth) - 3)],
+        "conv_stages": 3,
+    }, {
+        "embed_dims": [(2**i) * F for i in range(len(depth))],
+        "motion_dims": [0, 0, 0, 8 * F // depth[-2], 16 * F // depth[-1]],
+        "depths": depth,
+        "num_heads": [8 * (2**i) * F // 32 for i in range(len(depth) - 3)],
+        "window_sizes": [W, W],
+        "scales": [4 * (2**i) for i in range(len(depth) - 2)],
+        "hidden_dims": [4 * F for i in range(len(depth) - 3)],
+        "c": F,
+        "M": M,
+        "local_hidden_dims": 4 * F,
+        "local_num": 2,
+    }
+
+
+MODEL_CONFIG_DEFAULT = {
+    "LOGNAME": LOG,
+    "MODEL_TYPE": (mamba_extractor, mamba_estimation),
+    "MODEL_ARCH": init_model_config(F=16, depth=[2, 2, 2, 3, 3], M=False),
+}
+
+MODEL_CONFIG = {
+    "LOGNAME": LOG,
+    "MODEL_TYPE": (mamba_extractor, mamba_estimation),
+    "MODEL_ARCH": init_model_config(F=16, depth=[2, 2, 2, 3, 3], M=False),
+}

vfimamba_module/feature_extractor.py

@@ -0,0 +1,591 @@
+import math
+from typing import Callable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import repeat
+from mamba_ssm.ops.selective_scan_interface import selective_scan_fn
+from timm.models.layers import to_2tuple, trunc_normal_
+
+
+class ChannelAttention(nn.Module):
+    """Channel attention used in RCAN.
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        squeeze_factor (int): Channel squeeze factor. Default: 16.
+    """
+
+    def __init__(self, num_feat, squeeze_factor=16):
+        super(ChannelAttention, self).__init__()
+        self.attention = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(num_feat, num_feat // squeeze_factor, 1, padding=0),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(num_feat // squeeze_factor, num_feat, 1, padding=0),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        y = self.attention(x)
+        return x * y
+
+
+class CAB(nn.Module):
+    def __init__(
+        self, num_feat, is_light_sr=False, compress_ratio=3, squeeze_factor=30
+    ):
+        super(CAB, self).__init__()
+        self.cab = nn.Sequential(
+            nn.Conv2d(num_feat, num_feat // compress_ratio, 3, 1, 1),
+            nn.GELU(),
+            nn.Conv2d(num_feat // compress_ratio, num_feat, 3, 1, 1),
+            ChannelAttention(num_feat, squeeze_factor),
+        )
+
+    def forward(self, x):
+        return self.cab(x)
+
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(
+            in_planes,
+            out_planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=True,
+        ),
+        nn.PReLU(out_planes),
+    )
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class SS2D(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        d_state=16,
+        d_conv=3,
+        expand=2.0,
+        dt_rank="auto",
+        dt_min=0.001,
+        dt_max=0.1,
+        dt_init="random",
+        dt_scale=1.0,
+        dt_init_floor=1e-4,
+        dropout=0.0,
+        conv_bias=True,
+        bias=False,
+        device=None,
+        dtype=None,
+        **kwargs,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.d_model = d_model
+        self.d_state = d_state
+        self.d_conv = d_conv
+        self.expand = expand
+        self.d_inner = int(self.expand * self.d_model)
+        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank
+
+        self.in_proj = nn.Linear(
+            self.d_model, self.d_inner * 2, bias=bias, **factory_kwargs
+        )
+        self.conv2d = nn.Conv2d(
+            in_channels=self.d_inner,
+            out_channels=self.d_inner,
+            groups=self.d_inner,
+            bias=conv_bias,
+            kernel_size=d_conv,
+            padding=(d_conv - 1) // 2,
+            **factory_kwargs,
+        )
+        self.act = nn.SiLU()
+
+        self.x_proj = (
+            nn.Linear(
+                self.d_inner,
+                (self.dt_rank + self.d_state * 2),
+                bias=False,
+                **factory_kwargs,
+            ),
+            nn.Linear(
+                self.d_inner,
+                (self.dt_rank + self.d_state * 2),
+                bias=False,
+                **factory_kwargs,
+            ),
+            nn.Linear(
+                self.d_inner,
+                (self.dt_rank + self.d_state * 2),
+                bias=False,
+                **factory_kwargs,
+            ),
+            nn.Linear(
+                self.d_inner,
+                (self.dt_rank + self.d_state * 2),
+                bias=False,
+                **factory_kwargs,
+            ),
+        )
+        self.x_proj_weight = nn.Parameter(
+            torch.stack([t.weight for t in self.x_proj], dim=0)
+        )
+        del self.x_proj
+
+        self.dt_projs = (
+            self.dt_init(
+                self.dt_rank,
+                self.d_inner,
+                dt_scale,
+                dt_init,
+                dt_min,
+                dt_max,
+                dt_init_floor,
+                **factory_kwargs,
+            ),
+            self.dt_init(
+                self.dt_rank,
+                self.d_inner,
+                dt_scale,
+                dt_init,
+                dt_min,
+                dt_max,
+                dt_init_floor,
+                **factory_kwargs,
+            ),
+            self.dt_init(
+                self.dt_rank,
+                self.d_inner,
+                dt_scale,
+                dt_init,
+                dt_min,
+                dt_max,
+                dt_init_floor,
+                **factory_kwargs,
+            ),
+            self.dt_init(
+                self.dt_rank,
+                self.d_inner,
+                dt_scale,
+                dt_init,
+                dt_min,
+                dt_max,
+                dt_init_floor,
+                **factory_kwargs,
+            ),
+        )
+        self.dt_projs_weight = nn.Parameter(
+            torch.stack([t.weight for t in self.dt_projs], dim=0)
+        )
+        self.dt_projs_bias = nn.Parameter(
+            torch.stack([t.bias for t in self.dt_projs], dim=0)
+        )
+        del self.dt_projs
+
+        self.A_logs = self.A_log_init(self.d_state, self.d_inner, copies=4, merge=True)
+        self.Ds = self.D_init(self.d_inner, copies=4, merge=True)
+
+        self.selective_scan = selective_scan_fn
+
+        self.out_norm = nn.LayerNorm(self.d_inner)
+        self.out_proj = nn.Linear(
+            self.d_inner, self.d_model, bias=bias, **factory_kwargs
+        )
+        self.dropout = nn.Dropout(dropout) if dropout > 0.0 else None
+
+    @staticmethod
+    def dt_init(
+        dt_rank,
+        d_inner,
+        dt_scale=1.0,
+        dt_init="random",
+        dt_min=0.001,
+        dt_max=0.1,
+        dt_init_floor=1e-4,
+        **factory_kwargs,
+    ):
+        dt_proj = nn.Linear(dt_rank, d_inner, bias=True, **factory_kwargs)
+
+        dt_init_std = dt_rank**-0.5 * dt_scale
+        if dt_init == "constant":
+            nn.init.constant_(dt_proj.weight, dt_init_std)
+        elif dt_init == "random":
+            nn.init.uniform_(dt_proj.weight, -dt_init_std, dt_init_std)
+        else:
+            raise NotImplementedError
+
+        dt = torch.exp(
+            torch.rand(d_inner, **factory_kwargs)
+            * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min)
+        ).clamp(min=dt_init_floor)
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        with torch.no_grad():
+            dt_proj.bias.copy_(inv_dt)
+
+        dt_proj.bias._no_reinit = True
+
+        return dt_proj
+
+    @staticmethod
+    def A_log_init(d_state, d_inner, copies=1, device=None, merge=True):
+        A = repeat(
+            torch.arange(1, d_state + 1, dtype=torch.float32, device=device),
+            "n -> d n",
+            d=d_inner,
+        ).contiguous()
+        A_log = torch.log(A)
+        if copies > 1:
+            A_log = repeat(A_log, "d n -> r d n", r=copies)
+            if merge:
+                A_log = A_log.flatten(0, 1)
+        A_log = nn.Parameter(A_log)
+        A_log._no_weight_decay = True
+        return A_log
+
+    @staticmethod
+    def D_init(d_inner, copies=1, device=None, merge=True):
+        D = torch.ones(d_inner, device=device)
+        if copies > 1:
+            D = repeat(D, "n1 -> r n1", r=copies)
+            if merge:
+                D = D.flatten(0, 1)
+        D = nn.Parameter(D)
+        D._no_weight_decay = True
+        return D
+
+    def merge_x(self, x):
+        B, C, H, W = x.shape
+        L = 2 * H * W
+        x = x.view(B, -1, L // 2).transpose(1, 2)
+        x = torch.cat([x[: B // 2], x[B // 2 :]], dim=-1).reshape(B // 2, L, C)
+        return x.transpose(1, 2).contiguous()
+
+    def forward_core(self, x: torch.Tensor):
+        B, C, H, W = x.shape
+        L = 2 * H * W
+        K = 4
+        B = B // 2
+        x_hwwh = torch.stack(
+            [
+                self.merge_x(x),
+                self.merge_x(torch.transpose(x, dim0=2, dim1=3).contiguous()),
+            ],
+            dim=1,
+        ).view(B, 2, -1, L)
+        xs = torch.cat([x_hwwh, torch.flip(x_hwwh, dims=[-1])], dim=1)
+        x_dbl = torch.einsum(
+            "b k d l, k c d -> b k c l", xs.view(B, K, -1, L), self.x_proj_weight
+        )
+        dts, Bs, Cs = torch.split(
+            x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=2
+        )
+        dts = torch.einsum(
+            "b k r l, k d r -> b k d l", dts.view(B, K, -1, L), self.dt_projs_weight
+        )
+        xs = xs.float().view(B, -1, L)
+        dts = dts.contiguous().float().view(B, -1, L)
+        Bs = Bs.float().view(B, K, -1, L)
+        Cs = Cs.float().view(B, K, -1, L)
+        Ds = self.Ds.float().view(-1)
+        As = -torch.exp(self.A_logs.float()).view(-1, self.d_state)
+        dt_projs_bias = self.dt_projs_bias.float().view(-1)
+        out_y = self.selective_scan(
+            xs,
+            dts,
+            As,
+            Bs,
+            Cs,
+            Ds,
+            z=None,
+            delta_bias=dt_projs_bias,
+            delta_softplus=True,
+            return_last_state=False,
+        ).view(B, K, -1, L)
+        assert out_y.dtype == torch.float
+
+        inv_y = torch.flip(out_y[:, 2:4], dims=[-1]).view(B, 2, -1, L)
+
+        wh_y = (
+            torch.transpose(out_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3)
+            .contiguous()
+            .view(B, -1, L)
+        )
+        invwh_y = (
+            torch.transpose(inv_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3)
+            .contiguous()
+            .view(B, -1, L)
+        )
+
+        return out_y[:, 0], inv_y[:, 0], wh_y, invwh_y
+
+    def forward(self, x: torch.Tensor, **kwargs):
+        B, H, W, C = x.shape
+        xz = self.in_proj(x)
+        x, z = xz.chunk(2, dim=-1)
+        x = x.permute(0, 3, 1, 2).contiguous()
+        x = self.act(self.conv2d(x))
+        # print(x.shape)
+        y1, y2, y3, y4 = self.forward_core(x)
+        assert y1.dtype == torch.float32
+        y = y1 + y2 + y3 + y4
+        # print(y.shape)
+        y = (
+            torch.transpose(y, dim0=1, dim1=2)
+            .contiguous()
+            .view(B // 2, H * W, 2, int(self.expand * C))
+        )
+        y = torch.cat([y[:, :, 0], y[:, :, 1]], 0).view(
+            B, H, W, int(self.expand * C)
+        )  # .view(B//2, 2*H, 2*W, -1)
+
+        y = self.out_norm(y)
+        y = y * F.silu(z)
+        out = self.out_proj(y)
+
+        if self.dropout is not None:
+            out = self.dropout(out)
+        return out
+
+
+class VSSBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int = 0,
+        drop_path: float = 0,
+        norm_layer: Callable[..., torch.nn.Module] = nn.LayerNorm,
+        attn_drop_rate: float = 0,
+        d_state: int = 16,
+        mlp_ratio: float = 2.0,
+        **kwargs,
+    ):
+        super().__init__()
+        self.ln_1 = norm_layer(hidden_dim)
+        self.self_attention = SS2D(
+            d_model=hidden_dim,
+            d_state=d_state,
+            expand=mlp_ratio,
+            dropout=attn_drop_rate,
+            **kwargs,
+        )
+        self.skip_scale = nn.Parameter(torch.ones(hidden_dim))
+        self.conv_blk = CAB(hidden_dim)
+        self.ln_2 = nn.LayerNorm(hidden_dim)
+        self.skip_scale2 = nn.Parameter(torch.ones(hidden_dim))
+
+    def forward(self, input):
+        input = input.permute(0, 2, 3, 1).contiguous()
+        x = self.ln_1(input)
+
+        x = input * self.skip_scale + self.self_attention(x)
+        x = (
+            x * self.skip_scale2
+            + self.conv_blk(self.ln_2(x).permute(0, 3, 1, 2).contiguous())
+            .permute(0, 2, 3, 1)
+            .contiguous()
+        )
+
+        return x.permute(0, 3, 1, 2).contiguous()
+
+
+class BiMambaBlock(nn.Module):
+    def __init__(
+        self, dim, depth, norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False
+    ):
+        super().__init__()
+        self.dim = dim
+        self.depth = depth
+        self.blocks = nn.ModuleList()
+        for _ in range(depth):
+            self.blocks.append(
+                VSSBlock(
+                    hidden_dim=dim,
+                    norm_layer=nn.LayerNorm,
+                    d_state=16,
+                )
+            )
+
+    def forward(self, x):
+        for blk in self.blocks:
+            x = blk(x)
+        return x
+
+
+class OverlapPatchEmbed(nn.Module):
+    def __init__(self, patch_size=7, stride=4, in_chans=3, embed_dim=768):
+        super().__init__()
+        patch_size = to_2tuple(patch_size)
+
+        self.patch_size = patch_size
+        self.proj = nn.Conv2d(
+            in_chans,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=stride,
+            padding=(patch_size[0] // 2, patch_size[1] // 2),
+        )
+        self.norm = nn.LayerNorm(in_chans)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = x.permute(0, 2, 3, 1).contiguous()
+        x = self.norm(x).permute(0, 3, 1, 2).contiguous()
+        x = self.proj(x)
+        return x
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, depths=2, act_layer=nn.PReLU):
+        super().__init__()
+        layers = []
+        for i in range(depths):
+            if i == 0:
+                layers.append(nn.Conv2d(in_dim, out_dim, 3, 1, 1))
+            else:
+                layers.append(nn.Conv2d(out_dim, out_dim, 3, 1, 1))
+            layers.extend(
+                [
+                    act_layer(out_dim),
+                ]
+            )
+        self.conv = nn.Sequential(*layers)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class MambaFeature(nn.Module):
+    def __init__(
+        self,
+        img_size=64,
+        patch_size=1,
+        in_chans=3,
+        embed_dims=[16, 32, 64, 128, 256],
+        depths=(2, 2, 2, 2, 2),
+        conv_stages=2,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    ):
+        super(MambaFeature, self).__init__()
+
+        self.num_stages = len(embed_dims)
+
+        self.conv_stages = conv_stages
+
+        for i in range(self.num_stages):
+            if i == 0:
+                block = ConvBlock(in_chans, embed_dims[i], depths[i])
+            else:
+                if i < self.conv_stages:
+                    patch_embed = nn.Sequential(
+                        nn.Conv2d(embed_dims[i - 1], embed_dims[i], 3, 2, 1),
+                        nn.PReLU(embed_dims[i]),
+                    )
+                    block = ConvBlock(embed_dims[i], embed_dims[i], depths[i])
+                else:
+                    block = BiMambaBlock(embed_dims[i], depths[i])
+                    patch_embed = OverlapPatchEmbed(
+                        patch_size=3,
+                        stride=2,
+                        in_chans=embed_dims[i - 1],
+                        embed_dim=embed_dims[i],
+                    )
+                setattr(self, f"patch_embed{i}", patch_embed)
+            setattr(self, f"block{i}", block)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, x1, x2):
+        B = x1.shape[0]
+        x = torch.cat([x1, x2], 0)
+
+        features = []
+
+        for i in range(self.num_stages):
+            patch_embed = getattr(self, f"patch_embed{i}", None)
+            block = getattr(self, f"block{i}", None)
+
+            if i > 0:
+                x = patch_embed(x)
+
+            x = block(x)
+
+            features.append(x)
+
+        return features
+
+
+def feature_extractor(**kargs):
+    model = MambaFeature(**kargs)
+    return model

vfimamba_module/flow_estimation.py

@@ -0,0 +1,294 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .refine import Unet
+from .warplayer import warp
+
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(
+            in_planes,
+            out_planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=True,
+        ),
+        nn.PReLU(out_planes),
+    )
+
+
+class Head(nn.Module):
+    def __init__(self, in_planes, scale, c, in_else=17):
+        super(Head, self).__init__()
+        self.upsample = nn.Sequential(nn.PixelShuffle(2), nn.PixelShuffle(2))
+        self.scale = scale
+        self.conv = nn.Sequential(
+            conv(in_planes * 2 // (4 * 4) + in_else, c),
+            conv(c, c),
+            conv(c, 5),
+        )
+
+    def forward(self, motion_feature, x, flow):
+        motion_feature = self.upsample(motion_feature)
+        if self.scale != 4:
+            x = F.interpolate(
+                x, scale_factor=4.0 / self.scale, mode="bilinear", align_corners=False
+            )
+        if flow is not None:
+            if self.scale != 4:
+                flow = (
+                    F.interpolate(
+                        flow,
+                        scale_factor=4.0 / self.scale,
+                        mode="bilinear",
+                        align_corners=False,
+                    )
+                    * 4.0
+                    / self.scale
+                )
+            x = torch.cat((x, flow), 1)
+        x = self.conv(torch.cat([motion_feature, x], 1))
+        if self.scale != 4:
+            x = F.interpolate(
+                x, scale_factor=self.scale // 4, mode="bilinear", align_corners=False
+            )
+            flow = x[:, :4] * (self.scale // 4)
+        else:
+            flow = x[:, :4]
+        mask = x[:, 4:5]
+        return flow, mask
+
+
+class IFBlock(nn.Module):
+    def __init__(self, in_planes, c, scale):
+        super(IFBlock, self).__init__()
+        self.conv0 = nn.Sequential(
+            conv(in_planes, c // 2, 3, 2, 1),
+            conv(c // 2, c, 3, 2, 1),
+        )
+        self.convblock = nn.Sequential(
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+        )
+        self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)
+        self.scale = scale
+
+    def forward(self, x, flow):
+        scale = self.scale
+        if scale != 1:
+            x = F.interpolate(
+                x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
+            )
+            flow = (
+                F.interpolate(
+                    flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
+                )
+                * 1.0
+                / scale
+            )
+        x = torch.cat((x, flow), 1)
+        x = self.conv0(x)
+        x = self.convblock(x) + x
+        tmp = self.lastconv(x)
+        tmp = F.interpolate(
+            tmp, scale_factor=scale * 2, mode="bilinear", align_corners=False
+        )
+        flow = tmp[:, :4] * scale * 2
+        mask = tmp[:, 4:5]
+        return flow, mask
+
+
+class MultiScaleFlow(nn.Module):
+    def __init__(self, backbone, **kargs):
+        super(MultiScaleFlow, self).__init__()
+        self.flow_num_stage = len(kargs["hidden_dims"])
+        self.local_num = kargs["local_num"]
+        self.feature_bone = backbone
+        self.block = nn.ModuleList(
+            [
+                Head(
+                    kargs["embed_dims"][-1 - i],
+                    kargs["scales"][-1 - i],
+                    kargs["hidden_dims"][-1 - i],
+                    7 if i == 0 else 18,
+                )
+                for i in range(self.flow_num_stage)
+            ]
+        )
+        self.local_block = nn.ModuleList(
+            [
+                IFBlock(17, c=kargs["local_hidden_dims"], scale=2 - i)
+                for i in range(self.local_num)
+            ]
+        )
+        self.unet = Unet(kargs["c"] * 2, kargs["M"])
+
+    def warp_features(self, xs, flow):
+        y0 = []
+        y1 = []
+        B = xs[0].size(0) // 2
+        for x in xs:
+            y0.append(warp(x[:B], flow[:, 0:2]))
+            y1.append(warp(x[B:], flow[:, 2:4]))
+            flow = (
+                F.interpolate(
+                    flow,
+                    scale_factor=0.5,
+                    mode="bilinear",
+                    align_corners=False,
+                    recompute_scale_factor=False,
+                )
+                * 0.5
+            )
+        return y0, y1
+
+    def calculate_flow(self, imgs, timestep, local=False, af=None):
+        img0, img1 = imgs[:, :3], imgs[:, 3:6]
+        B = img0.size(0)
+        flow, mask = None, None
+        if af is None:
+            af = self.feature_bone(img0, img1)
+        timestep = (img0[:, :1].clone() * 0 + 1) * timestep
+        for i in range(self.flow_num_stage):
+            if flow is not None:
+                warped_img0 = warp(img0, flow[:, :2])
+                warped_img1 = warp(img1, flow[:, 2:4])
+                flow_, mask_ = self.block[i](
+                    torch.cat([af[-1 - i][:B], af[-1 - i][B:]], 1),
+                    torch.cat(
+                        (img0, img1, warped_img0, warped_img1, mask, timestep), 1
+                    ),
+                    flow,
+                )
+                flow = flow + flow_
+                mask = mask + mask_
+            else:
+                flow, mask = self.block[i](
+                    torch.cat([af[-1 - i][:B], af[-1 - i][B:]], 1),
+                    torch.cat((img0, img1, timestep), 1),
+                    None,
+                )
+
+        if local:
+            for i in range(self.local_num):
+                warped_img0 = warp(img0, flow[:, :2])
+                warped_img1 = warp(img1, flow[:, 2:4])
+
+                flow_d, mask_d = self.local_block[i](
+                    torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow
+                )
+                flow = flow + flow_d
+                mask = mask + mask_d
+
+        return flow, mask
+
+    def coraseWarp_and_Refine(self, imgs, af, flow, mask):
+        img0, img1 = imgs[:, :3], imgs[:, 3:6]
+        warped_img0 = warp(img0, flow[:, :2])
+        warped_img1 = warp(img1, flow[:, 2:4])
+        c0, c1 = self.warp_features(af, flow)
+        tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
+        res = tmp[:, :3] * 2 - 1
+        mask_ = torch.sigmoid(mask)
+        merged = warped_img0 * mask_ + warped_img1 * (1 - mask_)
+        pred = torch.clamp(merged + res, 0, 1)
+        return pred
+
+    def forward(self, x, local=False, timestep=0.5, scale=0):
+        if scale > 0:
+            x_o = x
+            x = F.interpolate(
+                x, scale_factor=scale, mode="bilinear", align_corners=False
+            )
+
+        img0, img1 = x[:, :3], x[:, 3:6]
+        B = x.size(0)
+        flow_list = []
+        merged = []
+        mask_list = []
+        warped_img0 = img0
+        warped_img1 = img1
+        flow = None
+        af = self.feature_bone(img0, img1)
+        timestep = (x[:, :1].clone() * 0 + 1) * (
+            timestep.float() if not isinstance(timestep, float) else timestep
+        )
+        for i in range(self.flow_num_stage):
+            if flow is not None:
+                flow_d, mask_d = self.block[i](
+                    torch.cat([af[-1 - i][:B], af[-1 - i][B:]], 1),
+                    torch.cat(
+                        (img0, img1, warped_img0, warped_img1, mask, timestep), 1  # noqa: F821
+                    ),
+                    flow,
+                )
+                flow = flow + flow_d
+                mask = mask + mask_d  # noqa: F821
+            else:
+                flow, mask = self.block[i](
+                    torch.cat([af[-1 - i][:B], af[-1 - i][B:]], 1),
+                    torch.cat((img0, img1, timestep), 1),
+                    None,
+                )
+            mask_list.append(torch.sigmoid(mask))
+            flow_list.append(flow)
+            warped_img0 = warp(img0, flow[:, :2])
+            warped_img1 = warp(img1, flow[:, 2:4])
+            merged.append(warped_img0 * mask_list[i] + warped_img1 * (1 - mask_list[i]))
+
+        if scale > 0:
+            img0, img1 = x_o[:, :3], x_o[:, 3:6]
+            af1 = self.feature_bone(img0, img1)
+            scale = img0.shape[3] / flow.shape[3]
+            flow = (
+                F.interpolate(
+                    flow, scale_factor=scale, mode="bilinear", align_corners=False
+                )
+                * scale
+            )
+            mask = F.interpolate(
+                mask, scale_factor=scale, mode="bilinear", align_corners=False
+            )
+            mask_ = torch.sigmoid(mask)
+            warped_img0 = warp(img0, flow[:, :2])
+            warped_img1 = warp(img1, flow[:, 2:4])
+            merged.append(warped_img0 * mask_ + warped_img1 * (1 - mask_))
+
+        if local:
+            flow_list = []
+            merged = []
+            mask_list = []
+
+            for i in range(self.local_num):
+                flow_d, mask_d = self.local_block[i](
+                    torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow
+                )
+                flow = flow + flow_d
+                mask = mask + mask_d
+
+                mask_list.append(torch.sigmoid(mask))
+                flow_list.append(flow)
+                warped_img0 = warp(img0, flow[:, :2])
+                warped_img1 = warp(img1, flow[:, 2:4])
+                merged.append(
+                    warped_img0 * mask_list[i] + warped_img1 * (1 - mask_list[i])
+                )
+
+        if scale:
+            c0, c1 = self.warp_features(af1, flow)
+        else:
+            c0, c1 = self.warp_features(af, flow)
+        tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
+        res = tmp[:, :3] * 2 - 1
+        pred = torch.clamp(merged[-1] + res, 0, 1)
+        return flow_list, mask_list, merged, pred

vfimamba_module/loss.py

@@ -0,0 +1,79 @@
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def gauss_kernel(channels=3):
+    kernel = torch.tensor(
+        [
+            [1.0, 4.0, 6.0, 4.0, 1],
+            [4.0, 16.0, 24.0, 16.0, 4.0],
+            [6.0, 24.0, 36.0, 24.0, 6.0],
+            [4.0, 16.0, 24.0, 16.0, 4.0],
+            [1.0, 4.0, 6.0, 4.0, 1.0],
+        ]
+    )
+    kernel /= 256.0
+    kernel = kernel.repeat(channels, 1, 1, 1)
+    kernel = kernel.to(device)
+    return kernel
+
+
+def downsample(x):
+    return x[:, :, ::2, ::2]
+
+
+def upsample(x):
+    cc = torch.cat(
+        [x, torch.zeros(x.shape[0], x.shape[1], x.shape[2], x.shape[3]).to(device)],
+        dim=3,
+    )
+    cc = cc.view(x.shape[0], x.shape[1], x.shape[2] * 2, x.shape[3])
+    cc = cc.permute(0, 1, 3, 2)
+    cc = torch.cat(
+        [
+            cc,
+            torch.zeros(x.shape[0], x.shape[1], x.shape[3], x.shape[2] * 2).to(device),
+        ],
+        dim=3,
+    )
+    cc = cc.view(x.shape[0], x.shape[1], x.shape[3] * 2, x.shape[2] * 2)
+    x_up = cc.permute(0, 1, 3, 2)
+    return conv_gauss(x_up, 4 * gauss_kernel(channels=x.shape[1]))
+
+
+def conv_gauss(img, kernel):
+    img = torch.nn.functional.pad(img, (2, 2, 2, 2), mode="reflect")
+    out = torch.nn.functional.conv2d(img, kernel, groups=img.shape[1])
+    return out
+
+
+def laplacian_pyramid(img, kernel, max_levels=3):
+    current = img
+    pyr = []
+    for level in range(max_levels):
+        filtered = conv_gauss(current, kernel)
+        down = downsample(filtered)
+        up = upsample(down)
+        diff = current - up
+        pyr.append(diff)
+        current = down
+    return pyr
+
+
+class LapLoss(torch.nn.Module):
+    def __init__(self, max_levels=5, channels=3):
+        super(LapLoss, self).__init__()
+        self.max_levels = max_levels
+        self.gauss_kernel = gauss_kernel(channels=channels)
+
+    def forward(self, input, target):
+        pyr_input = laplacian_pyramid(
+            img=input, kernel=self.gauss_kernel, max_levels=self.max_levels
+        )
+        pyr_target = laplacian_pyramid(
+            img=target, kernel=self.gauss_kernel, max_levels=self.max_levels
+        )
+        return sum(
+            torch.nn.functional.l1_loss(a, b) for a, b in zip(pyr_input, pyr_target)
+        )

vfimamba_module/padder.py

@@ -0,0 +1,24 @@
+import torch.nn.functional as F
+
+
+class InputPadder:
+    """Pads images such that dimensions are divisible by divisor"""
+
+    def __init__(self, dims, divisor=16):
+        self.ht, self.wd = dims[-2:]
+        pad_ht = (((self.ht // divisor) + 1) * divisor - self.ht) % divisor
+        pad_wd = (((self.wd // divisor) + 1) * divisor - self.wd) % divisor
+        self._pad = [
+            pad_wd // 2,
+            pad_wd - pad_wd // 2,
+            pad_ht // 2,
+            pad_ht - pad_ht // 2,
+        ]
+
+    def pad(self, *inputs):
+        return [F.pad(x, self._pad, mode="replicate") for x in inputs]
+
+    def unpad(self, x):
+        ht, wd = x.shape[-2:]
+        c = [self._pad[2], ht - self._pad[3], self._pad[0], wd - self._pad[1]]
+        return x[..., c[0] : c[1], c[2] : c[3]]

vfimamba_module/refine.py

@@ -0,0 +1,95 @@
+import math
+
+import torch
+import torch.nn as nn
+from timm.models.layers import trunc_normal_
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(
+            in_planes,
+            out_planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=True,
+        ),
+        nn.PReLU(out_planes),
+    )
+
+
+def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
+    return nn.Sequential(
+        torch.nn.ConvTranspose2d(
+            in_channels=in_planes,
+            out_channels=out_planes,
+            kernel_size=4,
+            stride=2,
+            padding=1,
+            bias=True,
+        ),
+        nn.PReLU(out_planes),
+    )
+
+
+class Conv2(nn.Module):
+    def __init__(self, in_planes, out_planes, stride=2):
+        super(Conv2, self).__init__()
+        self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
+        self.conv2 = conv(out_planes, out_planes, 3, 1, 1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+
+
+class Unet(nn.Module):
+    def __init__(self, c, M=False, out=3):
+        super(Unet, self).__init__()
+        self.down0 = Conv2(17 + c, 2 * c)
+        self.down1 = Conv2(4 * c, 4 * c)
+        self.down2 = Conv2(8 * c, 8 * c)
+        self.down3 = Conv2(16 * c, 16 * c)
+        self.up0 = deconv(32 * c if not M else 16 * c, 8 * c)
+        self.up1 = deconv(16 * c, 4 * c)
+        self.up2 = deconv(8 * c, 2 * c)
+        self.up3 = deconv(4 * c, c)
+        self.conv = nn.Conv2d(c, out, 3, 1, 1)
+        self.M = M
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1):
+        s0 = self.down0(
+            torch.cat(
+                (img0, img1, warped_img0, warped_img1, mask, flow, c0[0], c1[0]), 1
+            )
+        )
+        s1 = self.down1(torch.cat((s0, c0[1], c1[1]), 1))
+        s2 = self.down2(torch.cat((s1, c0[2], c1[2]), 1))
+        s3 = self.down3(torch.cat((s2, c0[3], c1[3]), 1))
+        x = self.up0(torch.cat((s3, c0[4], c1[4]), 1) if not self.M else s3)
+        x = self.up1(torch.cat((x, s2), 1))
+        x = self.up2(torch.cat((x, s1), 1))
+        x = self.up3(torch.cat((x, s0), 1))
+        x = self.conv(x)
+        return torch.sigmoid(x)

vfimamba_module/vfi_mamba_model.py

@@ -0,0 +1,89 @@
+import math
+
+import torch
+import tqdm
+from comfy.utils import ProgressBar
+
+from . import config as cfg
+from .padder import InputPadder
+from .Trainer_finetune import Model
+
+
+def run_vfi_mamba(
+    src_video: torch.Tensor,
+    model_name: str,
+    model_path: str,
+    vfi_scale=2,
+    scale: float = 0.0,
+):
+    print(f"{model_name=}")
+    print(f"{model_path=}")
+    TTA = False
+    cfg.MODEL_CONFIG["LOGNAME"] = model_name
+    if model_name == "VFIMamba":
+        TTA = True
+        cfg.MODEL_CONFIG["MODEL_ARCH"] = cfg.init_model_config(
+            F=32, depth=[2, 2, 2, 3, 3]
+        )
+    else:
+        cfg.MODEL_CONFIG = cfg.MODEL_CONFIG_DEFAULT.copy()
+    print()
+    model = Model(-1, cfg.MODEL_CONFIG)
+    model.load_model(model_path=model_path)
+    model.eval()
+    model.device()
+
+    src_video = src_video.permute(0, 3, 1, 2).to("cuda")
+
+    out_frames: list[torch.Tensor] = []
+
+    pbar = ProgressBar(src_video.shape[0])
+    for i in tqdm.tqdm(range(src_video.shape[0] - 1), "Generating frames"):
+        frame_a = src_video[i].unsqueeze(dim=0)
+        frame_b = src_video[i + 1].unsqueeze(dim=0)
+
+        padder = InputPadder(frame_a.shape, divisor=32)
+        I0_, I2_ = padder.pad(frame_a, frame_b)
+
+        if vfi_scale == 2:
+            out = model.inference(I0_, I2_, True, TTA=TTA, fast_TTA=TTA, scale=scale)
+            mid_frame = padder.unpad(out)
+            out_frames.append(frame_a)
+            out_frames.append(mid_frame)
+        else:
+            frames = _recursive_generator(
+                model=model,
+                TTA=TTA,
+                scale=scale,
+                frame1=I0_,
+                frame2=I2_,
+                num_recursions=int(math.log2(vfi_scale)),
+                index=vfi_scale // 2,
+            )
+            frames = sorted(frames, key=lambda x: x[1])
+            for pred, _ in frames:
+                out_frames.append(padder.unpad(pred))
+        pbar.update_absolute(i, src_video.shape[0])
+
+    out_frames.append(src_video[-1].unsqueeze(0))
+
+    out_tensor = torch.cat(out_frames, dim=0).permute(0, 2, 3, 1)
+    print(f"{out_tensor.shape=}")
+    return out_tensor
+
+
+def _recursive_generator(
+    model, TTA, scale, frame1: torch.Tensor, frame2: torch.Tensor, num_recursions, index
+):
+    if num_recursions == 0:
+        return [(frame1, index)]
+    else:
+        mid_frame = model.inference(
+            frame1, frame2, True, TTA=TTA, fast_TTA=TTA, scale=scale
+        )
+        id = 2 ** (num_recursions - 1)
+        return _recursive_generator(
+            model, TTA, scale, frame1, mid_frame, num_recursions - 1, index - id
+        ) + _recursive_generator(
+            model, TTA, scale, mid_frame, frame2, num_recursions - 1, index + id
+        )

vfimamba_module/warplayer.py

@@ -0,0 +1,37 @@
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+backwarp_tenGrid = {}
+
+
+def warp(tenInput, tenFlow):
+    k = (str(tenFlow.device), str(tenFlow.size()))
+    if k not in backwarp_tenGrid:
+        tenHorizontal = (
+            torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device)
+            .view(1, 1, 1, tenFlow.shape[3])
+            .expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1)
+        )
+        tenVertical = (
+            torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device)
+            .view(1, 1, tenFlow.shape[2], 1)
+            .expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3])
+        )
+        backwarp_tenGrid[k] = torch.cat([tenHorizontal, tenVertical], 1).to(device)
+
+    tenFlow = torch.cat(
+        [
+            tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0),
+            tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0),
+        ],
+        1,
+    )
+
+    g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1)
+    return torch.nn.functional.grid_sample(
+        input=tenInput,
+        grid=g,
+        mode="bilinear",
+        padding_mode="border",
+        align_corners=True,
+    )

workflow/workflow_vfi_mamba.json

@@ -0,0 +1,316 @@
+{
+  "last_node_id": 9,
+  "last_link_id": 13,
+  "nodes": [
+    {
+      "id": 2,
+      "type": "VHS_LoadVideo",
+      "pos": [
+        393,
+        210
+      ],
+      "size": [
+        235.1999969482422,
+        398.1134866403979
+      ],
+      "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": [
+            12
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "links": [
+            7
+          ],
+          "shape": 3,
+          "slot_index": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_LoadVideo"
+      },
+      "widgets_values": {
+        "video": "cat.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.mp4",
+            "type": "input",
+            "format": "video/mp4",
+            "select_every_nth": 1
+          }
+        }
+      }
+    },
+    {
+      "id": 4,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        1052,
+        219
+      ],
+      "size": [
+        315,
+        487
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 13
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        },
+        {
+          "name": "frame_rate",
+          "type": "FLOAT",
+          "link": 8,
+          "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": 20,
+        "loop_count": 0,
+        "filename_prefix": "VFIMamba",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 0,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": true,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "VFIMamba_00011.mp4",
+            "subfolder": "",
+            "type": "output",
+            "format": "video/h264-mp4",
+            "frame_rate": 20
+          }
+        }
+      }
+    },
+    {
+      "id": 7,
+      "type": "VHS_VideoInfoLoaded",
+      "pos": [
+        680,
+        360
+      ],
+      "size": {
+        "0": 304.79998779296875,
+        "1": 106
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "link": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "fps🟦",
+          "type": "FLOAT",
+          "links": [
+            8
+          ],
+          "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": 9,
+      "type": "VFI_MAMBA",
+      "pos": [
+        691,
+        205
+      ],
+      "size": {
+        "0": 283.76104736328125,
+        "1": 106
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "src_video",
+          "type": "IMAGE",
+          "link": 12
+        }
+      ],
+      "outputs": [
+        {
+          "name": "res_video",
+          "type": "IMAGE",
+          "links": [
+            13
+          ],
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VFI_MAMBA"
+      },
+      "widgets_values": [
+        5,
+        "VFIMamba_S"
+      ]
+    }
+  ],
+  "links": [
+    [
+      7,
+      2,
+      3,
+      7,
+      0,
+      "VHS_VIDEOINFO"
+    ],
+    [
+      8,
+      7,
+      0,
+      4,
+      4,
+      "FLOAT"
+    ],
+    [
+      12,
+      2,
+      0,
+      9,
+      0,
+      "IMAGE"
+    ],
+    [
+      13,
+      9,
+      0,
+      4,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1.4641000000000044,
+      "offset": [
+        -294.3910191805387,
+        -112.89932302330475
+      ]
+    }
+  },
+  "version": 0.4
+}