diff --git a/.gitignore b/.gitignore index dee2e10..10f3a62 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ *.pth +*.pkl # Byte-compiled / optimized / DLL files __pycache__ diff --git a/README.md b/README.md index 0ef2865..cd0e1b3 100644 --- a/README.md +++ b/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) \ No newline at end of file diff --git a/__init__.py b/__init__.py index 2bdd7e6..834c8f2 100644 --- a/__init__.py +++ b/__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"] \ No newline at end of file diff --git a/constants.py b/constants.py index 8f31860..2ee5fb7 100644 --- a/constants.py +++ b/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" \ No newline at end of file +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" diff --git a/example_vfi_mamba.png b/example_vfi_mamba.png new file mode 100644 index 0000000..95d1608 Binary files /dev/null and b/example_vfi_mamba.png differ diff --git a/pyproject.toml b/pyproject.toml index 7ee4ca9..f74f332 100644 --- a/pyproject.toml +++ b/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"] diff --git a/vfi_mamba_run.py b/vfi_mamba_run.py new file mode 100644 index 0000000..6613ac4 --- /dev/null +++ b/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, + ), + ) diff --git a/vfimamba_module/Trainer_finetune.py b/vfimamba_module/Trainer_finetune.py new file mode 100644 index 0000000..e4004eb --- /dev/null +++ b/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 diff --git a/vfimamba_module/config.py b/vfimamba_module/config.py new file mode 100644 index 0000000..7d1b659 --- /dev/null +++ b/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), +} diff --git a/vfimamba_module/feature_extractor.py b/vfimamba_module/feature_extractor.py new file mode 100644 index 0000000..444b8d0 --- /dev/null +++ b/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 diff --git a/vfimamba_module/flow_estimation.py b/vfimamba_module/flow_estimation.py new file mode 100644 index 0000000..7f1a62f --- /dev/null +++ b/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 diff --git a/vfimamba_module/loss.py b/vfimamba_module/loss.py new file mode 100644 index 0000000..2cdd990 --- /dev/null +++ b/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) + ) diff --git a/vfimamba_module/padder.py b/vfimamba_module/padder.py new file mode 100644 index 0000000..4b0cd28 --- /dev/null +++ b/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]] diff --git a/vfimamba_module/refine.py b/vfimamba_module/refine.py new file mode 100644 index 0000000..9da51b4 --- /dev/null +++ b/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) diff --git a/vfimamba_module/vfi_mamba_model.py b/vfimamba_module/vfi_mamba_model.py new file mode 100644 index 0000000..1010590 --- /dev/null +++ b/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 + ) diff --git a/vfimamba_module/warplayer.py b/vfimamba_module/warplayer.py new file mode 100644 index 0000000..4a95aee --- /dev/null +++ b/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, + ) diff --git a/workflow/workflow_vfi_mamba.json b/workflow/workflow_vfi_mamba.json new file mode 100644 index 0000000..a621712 --- /dev/null +++ b/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 +} \ No newline at end of file