diff --git a/README.md b/README.md new file mode 100644 index 0000000..b63cfb6 --- /dev/null +++ b/README.md @@ -0,0 +1,64 @@ +# ComfyUI-Lee-RIFE + +Fast, **GPU-batched** RIFE frame interpolation for ComfyUI — a self-contained node +with the weights bundled, no extra dependencies. + +## What it does + +* Runs RIFE (`rife49.pth`, architecture v4.7) **on the GPU in batches** instead of +one frame-pair at a time. +* Enables RIFE's **real bidirectional ensemble** (a forward and a time-reversed +pass, averaged) for cleaner motion. +* `fp16` by default (visually identical to `fp32` for interpolation, roughly 2× +faster on modern NVIDIA cards); `fp32` selectable. +* Automatically halves the batch on out-of-memory, so it scales from small cards +to big ones. + +## Requirements + +* ComfyUI +* An NVIDIA GPU with a CUDA build of PyTorch (the same one ComfyUI already uses). +* **No extra pip installs** — it only uses what ComfyUI provides. + +## Install + +1. Copy the `ComfyUI-Lee-RIFE` folder into `ComfyUI/custom_nodes/` +(or `git clone` it there). +2. Make sure `models/rife49.pth` is present (it ships with this repo; +see `models/HOW-TO-ADD-WEIGHTS.txt` if it is missing). +3. Restart ComfyUI completely (the process, not just the browser tab). + +The node appears as **“RIFE GPU Fast (Lee)”** under the **Lee/RIFE** category. + +## Usage + +Wire your frames into `frames`. For doubling 15 fps → 30 fps, set `multiplier = 2` +and set your video-output node to 30 fps. + +|Parameter|Meaning| +|-|-| +|`multiplier`|Output frames per input frame (2 = double the frame rate).| +|`ensemble`|Bidirectional pass (better quality, \~2× the work). On by default.| +|`scale\_factor`|Internal flow scale; leave at `1.0` unless resolution is very high.| +|`precision`|`fp16` (fast) or `fp32` (bit-faithful to the original node).| +|`batch\_size`|Frame-pairs per GPU batch. Higher = better GPU use; lowers itself on OOM.| + +**Output count:** for `multiplier = m` and `N` input frames you get `m × N` +frames (e.g. 81 frames at 15 fps → 162 frames at 30 fps, same duration). + +On the first run you will see a console line like +`\[ComfyUI-Lee-RIFE] loaded rife49.pth (arch 4.7) on cuda:0 / torch.float16`, +confirming it is on the GPU. + +## Credits \& license + +This node's own code is **MIT** (see `LICENSE`), © 2026 Smai-Lee. + +It builds on: + +* **RIFE / Practical-RIFE** (Zhewei Huang et al.) — the model and architecture (MIT). +* **ComfyUI-Frame-Interpolation** (Fannovel16) — `rife\_arch.py` is vendored from it (MIT). + +Full attribution and the upstream license texts are in **`THIRD-PARTY-NOTICES.md`**. +Please cite the RIFE paper if you use this in research (citation in the notices file). + diff --git a/THIRD-PARTY-NOTICES.md b/THIRD-PARTY-NOTICES.md new file mode 100644 index 0000000..fd85bdf --- /dev/null +++ b/THIRD-PARTY-NOTICES.md @@ -0,0 +1,96 @@ +# Third-Party Notices + +ComfyUI-Lee-RIFE's own code is licensed under the MIT License (see `LICENSE`). +It bundles and builds upon the third-party components listed below. Each is used +under its own license; the original copyright and license notices are reproduced +here as required. + +\--- + +## 1\) `rife\_arch.py` — RIFE architecture (IFNet) implementation + +Vendored **verbatim** from **ComfyUI-Frame-Interpolation** by Fannovel16: +https://github.com/Fannovel16/ComfyUI-Frame-Interpolation + +That file in turn derives from RIFE / Practical-RIFE and vs-rife — see the +attribution header inside `rife\_arch.py`: + +* https://github.com/hzwer/Practical-RIFE +* https://github.com/HolyWu/vs-rife + +License: **MIT** + +``` +MIT License + +Copyright (c) 2023 Fannovel16 + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +``` + +\--- + +## 2\) Model weights — `models/rife47/49.pth` (RIFE v4.7/v4.9) + +Source: **RIFE / Practical-RIFE** by Zhewei Huang (hzwer) and contributors: + +* https://github.com/hzwer/Practical-RIFE +* https://github.com/megvii-research/ECCV2022-RIFE + +The RIFE models are released by their authors under the **MIT License** +(the Practical-RIFE README states the model files are distributed under the +same MIT license as the project). © the RIFE / Practical-RIFE authors. For the +authoritative copyright line, see the `LICENSE` file in the repositories above. + +MIT permission terms (same text as above) apply to these weights: + +``` +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND ... +``` + +If you use this in academic work, please cite RIFE: + +``` +@inproceedings{huang2022rife, + title = {Real-Time Intermediate Flow Estimation for Video Frame Interpolation}, + author = {Huang, Zhewei and Zhang, Tianyuan and Heng, Wen and Shi, Boxin and Zhou, Shuchang}, + booktitle = {European Conference on Computer Vision (ECCV)}, + year = {2022} +} +``` + +\--- + +## 3\) ComfyUI (host application) + +This is a custom node for ComfyUI (https://github.com/comfyanonymous/ComfyUI), +which is licensed under GPL-3.0. ComfyUI is **not** bundled or redistributed with +this package — it is the host that loads this node at runtime. This node's own +code is provided under the MIT License (see `LICENSE`). + diff --git a/models/HOW-TO-ADD-WEIGHTS.txt b/models/HOW-TO-ADD-WEIGHTS.txt new file mode 100644 index 0000000..3797436 --- /dev/null +++ b/models/HOW-TO-ADD-WEIGHTS.txt @@ -0,0 +1,13 @@ +This folder must contain the RIFE weights file: rife49.pth + +This repository is meant to ship with rife49.pth already here. If it is missing +(e.g. you obtained a copy without the binary), copy rife49.pth into this folder. + +Where to get it: + - From an existing ComfyUI-Frame-Interpolation install: + ComfyUI/custom_nodes/comfyui-frame-interpolation/ckpts/rife/rife49.pth + - From RIFE / Practical-RIFE: + https://github.com/hzwer/Practical-RIFE + +rife47.pth is also supported (same architecture) — just drop it here too. +The file is about 21 MB. No Git LFS is required. diff --git a/models/rife49.pth b/models/rife49.pth new file mode 100644 index 0000000..c4056e4 Binary files /dev/null and b/models/rife49.pth differ diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..32c4871 --- /dev/null +++ b/nodes.py @@ -0,0 +1,164 @@ +# nodes.py -- ComfyUI-Lee-RIFE +# +# GPU-batched RIFE frame interpolation as a self-contained ComfyUI node. +# +# Why batched: the common RIFE node pushes frame-pairs through ONE AT A TIME, +# copies each pair CPU<->GPU and flushes the CUDA cache every few frames. On a +# fast card the RIFE forward is tiny, so the GPU idles between those serial +# steps (low GPU utilisation). This node batches the pairs and skips the +# per-frame flush, keeping the GPU busy. It also enables RIFE's real +# bidirectional "ensemble" (averaging a forward and a time-reversed pass). +# +# The architecture (rife_arch.py) and the weights (models/rife49.pth) come from +# RIFE / Practical-RIFE and ComfyUI-Frame-Interpolation, both MIT-licensed. +# See THIRD-PARTY-NOTICES.md. + +import os +import torch + +from .rife_arch import IFNet + +try: + from comfy.model_management import get_torch_device, soft_empty_cache +except Exception: # allows importing/testing outside ComfyUI + def get_torch_device(): + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + def soft_empty_cache(*a, **k): + if torch.cuda.is_available(): + torch.cuda.empty_cache() + +HERE = os.path.dirname(os.path.abspath(__file__)) +MODELS_DIR = os.path.join(HERE, "models") + +# RIFE checkpoint -> architecture version +CKPT_ARCH = { + "rife40.pth": "4.0", "rife41.pth": "4.0", "rife42.pth": "4.2", + "rife43.pth": "4.3", "rife44.pth": "4.3", "rife45.pth": "4.5", + "rife46.pth": "4.6", "rife47.pth": "4.7", "rife48.pth": "4.7", + "rife49.pth": "4.7", +} + + +def _available_ckpts(): + found = [n for n in CKPT_ARCH if os.path.isfile(os.path.join(MODELS_DIR, n))] + # Fallback so the node still registers even before the weights are in place. + return found or ["rife49.pth"] + + +class LeeRIFEGPU: + _cache = {} # (ckpt_name, dtype_str) -> loaded model + + @classmethod + def INPUT_TYPES(s): + ckpts = _available_ckpts() + default = "rife49.pth" if "rife49.pth" in ckpts else ckpts[0] + return { + "required": { + "frames": ("IMAGE",), + "ckpt_name": (ckpts, {"default": default}), + "multiplier": ("INT", {"default": 2, "min": 1, "max": 16}), + "ensemble": ("BOOLEAN", {"default": True}), + "scale_factor": ([0.25, 0.5, 1.0, 2.0, 4.0], {"default": 1.0}), + "precision": (["fp16", "fp32"], {"default": "fp16"}), + "batch_size": ("INT", {"default": 16, "min": 1, "max": 256}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "interpolate" + CATEGORY = "Lee/RIFE" + + def _get_model(self, ckpt_name, dtype): + key = (ckpt_name, str(dtype)) + if key in self._cache: + return self._cache[key] + path = os.path.join(MODELS_DIR, ckpt_name) + if not os.path.isfile(path): + raise FileNotFoundError( + f"{ckpt_name} not found in {MODELS_DIR}. " + f"Place the RIFE weights file there (see models/HOW-TO-ADD-WEIGHTS.txt)." + ) + arch = CKPT_ARCH.get(ckpt_name, "4.7") + model = IFNet(arch_ver=arch) + try: + sd = torch.load(path, map_location="cpu", weights_only=True) + except Exception: + sd = torch.load(path, map_location="cpu", weights_only=False) + if isinstance(sd, dict) and "state_dict" in sd: + sd = sd["state_dict"] + model.load_state_dict(sd) + model.eval().to(get_torch_device()) + if dtype == torch.float16: + model.half() + print(f"[ComfyUI-Lee-RIFE] loaded {ckpt_name} (arch {arch}) on " + f"{get_torch_device()} / {dtype}") + self._cache[key] = model + return model + + @torch.no_grad() + def interpolate(self, frames, ckpt_name, multiplier, ensemble, + scale_factor, precision, batch_size): + device = get_torch_device() + dtype = torch.float16 if precision == "fp16" else torch.float32 + model = self._get_model(ckpt_name, dtype) + + N = frames.shape[0] + if N < 2 or multiplier < 2: + return (frames,) + + # ComfyUI IMAGE is (N,H,W,C) in [0,1]; RIFE wants (N,C,H,W). + # Stays on CPU here; only the per-batch chunks are moved to the GPU. + seq = frames.permute(0, 3, 1, 2).contiguous() + img0_all = seq[:-1] + img1_all = seq[1:] + P = img0_all.shape[0] + scale_list = [8 / scale_factor, 4 / scale_factor, + 2 / scale_factor, 1 / scale_factor] + + # One batched pass per intermediate timestep k/multiplier. + mids = [] # list over k, each tensor (P,C,H,W) on CPU / float32 + for k in range(1, multiplier): + t = k / multiplier + parts = [] + s = 0 + chunk = int(batch_size) + while s < P: + a = img0_all[s:s + chunk].to(device, dtype) + b = img1_all[s:s + chunk].to(device, dtype) + try: + r = model(a, b, timestep=t, scale_list=list(scale_list), + training=False, fastmode=True, ensemble=ensemble) + except RuntimeError as e: + del a, b + if "out of memory" not in str(e).lower(): + raise + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if chunk == 1: + raise + chunk = max(1, chunk // 2) # auto-shrink and retry + continue + parts.append(r.float().clamp(0, 1).cpu()) + del a, b, r + s += chunk + mids.append(torch.cat(parts, 0)) + + # Interleave originals + midpoints, pad to multiplier*N output frames. + seq_cpu = seq.float().cpu() + out = [] + for i in range(N - 1): + out.append(seq_cpu[i]) + for k in range(multiplier - 1): + out.append(mids[k][i]) + out.append(seq_cpu[N - 1]) + while len(out) < multiplier * N: + out.append(seq_cpu[N - 1]) + + result = torch.stack(out, 0).permute(0, 2, 3, 1).contiguous() # (M*N,H,W,C) + soft_empty_cache() + return (result[..., :3],) + + +NODE_CLASS_MAPPINGS = {"LeeRIFEGPU": LeeRIFEGPU} +NODE_DISPLAY_NAME_MAPPINGS = {"LeeRIFEGPU": "RIFE GPU Fast (Lee)"} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d898c78 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +# ComfyUI-Lee-RIFE has NO extra Python dependencies. +# +# PyTorch is provided by your ComfyUI installation. Do NOT let a package manager +# reinstall torch over it — that can replace your CUDA-enabled build and break +# GPU support. This file is intentionally empty (comments only). diff --git a/rife_arch.py b/rife_arch.py new file mode 100644 index 0000000..cf2b454 --- /dev/null +++ b/rife_arch.py @@ -0,0 +1,586 @@ +""" +26-Dez-21 +https://github.com/hzwer/Practical-RIFE +https://github.com/hzwer/Practical-RIFE/blob/main/model/warplayer.py +https://github.com/HolyWu/vs-rife/blob/master/vsrife/__init__.py +""" +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import AdamW +import torch +import torch.nn.functional as F +import torch.nn as nn +import torch.optim as optim +import warnings +from comfy.model_management import get_torch_device + +device = get_torch_device() +backwarp_tenGrid = {} + + +class ResConv(nn.Module): + def __init__(self, c, dilation=1): + super(ResConv, self).__init__() + self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1) + self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x): + return self.relu(self.conv(x) * self.beta + x) + + +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) + + if tenInput.type() == "torch.cuda.HalfTensor": + g = g.half() + + padding_mode = "border" + if device.type == "mps": + # https://github.com/pytorch/pytorch/issues/125098 + padding_mode = "zeros" + g = g.clamp(-1, 1) + return torch.nn.functional.grid_sample( + input=tenInput, + grid=g, + mode="bilinear", + padding_mode=padding_mode, + align_corners=True, + ) + + +def conv( + in_planes, + out_planes, + kernel_size=3, + stride=1, + padding=1, + dilation=1, + arch_ver="4.0", +): + if arch_ver == "4.0": + 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), + ) + if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]: + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.LeakyReLU(0.2, True), + ) + + +def conv_woact(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, + ), + ) + + +def conv_woact(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, + ) + ) + + +def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1, arch_ver="4.0"): + if arch_ver == "4.0": + 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), + ) + if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]: + return nn.Sequential( + torch.nn.ConvTranspose2d( + in_channels=in_planes, + out_channels=out_planes, + kernel_size=4, + stride=2, + padding=1, + bias=True, + ), + nn.LeakyReLU(0.2, True), + ) + + +class Conv2(nn.Module): + def __init__(self, in_planes, out_planes, stride=2, arch_ver="4.0"): + super(Conv2, self).__init__() + self.conv1 = conv(in_planes, out_planes, 3, stride, 1, arch_ver=arch_ver) + self.conv2 = conv(out_planes, out_planes, 3, 1, 1, arch_ver=arch_ver) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + +class IFBlock(nn.Module): + def __init__(self, in_planes, c=64, arch_ver="4.0"): + super(IFBlock, self).__init__() + self.arch_ver = arch_ver + self.conv0 = nn.Sequential( + conv(in_planes, c // 2, 3, 2, 1, arch_ver=arch_ver), + conv(c // 2, c, 3, 2, 1, arch_ver=arch_ver), + ) + self.arch_ver = arch_ver + + if arch_ver in ["4.0", "4.2", "4.3"]: + self.convblock = nn.Sequential( + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + ) + self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1) + + if arch_ver in ["4.5", "4.6", "4.7", "4.10"]: + self.convblock = nn.Sequential( + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ) + if arch_ver == "4.5": + self.lastconv = nn.Sequential( + nn.ConvTranspose2d(c, 4 * 5, 4, 2, 1), nn.PixelShuffle(2) + ) + if arch_ver in ["4.6", "4.7", "4.10"]: + self.lastconv = nn.Sequential( + nn.ConvTranspose2d(c, 4 * 6, 4, 2, 1), nn.PixelShuffle(2) + ) + + def forward(self, x, flow=None, scale=1): + x = F.interpolate( + x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False + ) + if flow is not None: + flow = ( + F.interpolate( + flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False + ) + * 1.0 + / scale + ) + x = torch.cat((x, flow), 1) + feat = self.conv0(x) + if self.arch_ver == "4.0": + feat = self.convblock(feat) + feat + if self.arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]: + feat = self.convblock(feat) + + tmp = self.lastconv(feat) + if self.arch_ver in ["4.0", "4.2", "4.3"]: + tmp = F.interpolate( + tmp, scale_factor=scale * 2, mode="bilinear", align_corners=False + ) + flow = tmp[:, :4] * scale * 2 + if self.arch_ver in ["4.5", "4.6", "4.7", "4.10"]: + tmp = F.interpolate( + tmp, scale_factor=scale, mode="bilinear", align_corners=False + ) + flow = tmp[:, :4] * scale + mask = tmp[:, 4:5] + return flow, mask + + +class Contextnet(nn.Module): + def __init__(self, arch_ver="4.0"): + super(Contextnet, self).__init__() + c = 16 + self.conv1 = Conv2(3, c, arch_ver=arch_ver) + self.conv2 = Conv2(c, 2 * c, arch_ver=arch_ver) + self.conv3 = Conv2(2 * c, 4 * c, arch_ver=arch_ver) + self.conv4 = Conv2(4 * c, 8 * c, arch_ver=arch_ver) + + def forward(self, x, flow): + x = self.conv1(x) + flow = ( + F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) + * 0.5 + ) + f1 = warp(x, flow) + x = self.conv2(x) + flow = ( + F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) + * 0.5 + ) + f2 = warp(x, flow) + x = self.conv3(x) + flow = ( + F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) + * 0.5 + ) + f3 = warp(x, flow) + x = self.conv4(x) + flow = ( + F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) + * 0.5 + ) + f4 = warp(x, flow) + return [f1, f2, f3, f4] + + +class Unet(nn.Module): + def __init__(self, arch_ver="4.0"): + super(Unet, self).__init__() + c = 16 + self.down0 = Conv2(17, 2 * c, arch_ver=arch_ver) + self.down1 = Conv2(4 * c, 4 * c, arch_ver=arch_ver) + self.down2 = Conv2(8 * c, 8 * c, arch_ver=arch_ver) + self.down3 = Conv2(16 * c, 16 * c, arch_ver=arch_ver) + self.up0 = deconv(32 * c, 8 * c, arch_ver=arch_ver) + self.up1 = deconv(16 * c, 4 * c, arch_ver=arch_ver) + self.up2 = deconv(8 * c, 2 * c, arch_ver=arch_ver) + self.up3 = deconv(4 * c, c, arch_ver=arch_ver) + self.conv = nn.Conv2d(c, 3, 3, 1, 1) + + 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), 1) + ) + s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1)) + s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1)) + s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1)) + x = self.up0(torch.cat((s3, c0[3], c1[3]), 1)) + 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) + + +""" +currently supports 4.0-4.12 + +4.0: 4.0, 4.1 +4.2: 4.2 +4.3: 4.3, 4.4 +4.5: 4.5 +4.6: 4.6 +4.7: 4.7, 4.8, 4.9 +4.10: 4.10 4.11 4.12 +""" + + +class IFNet(nn.Module): + def __init__(self, arch_ver="4.0"): + super(IFNet, self).__init__() + self.arch_ver = arch_ver + if arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]: + self.block0 = IFBlock(7, c=192, arch_ver=arch_ver) + self.block1 = IFBlock(8 + 4, c=128, arch_ver=arch_ver) + self.block2 = IFBlock(8 + 4, c=96, arch_ver=arch_ver) + self.block3 = IFBlock(8 + 4, c=64, arch_ver=arch_ver) + if arch_ver in ["4.7"]: + self.block0 = IFBlock(7 + 8, c=192, arch_ver=arch_ver) + self.block1 = IFBlock(8 + 4 + 8, c=128, arch_ver=arch_ver) + self.block2 = IFBlock(8 + 4 + 8, c=96, arch_ver=arch_ver) + self.block3 = IFBlock(8 + 4 + 8, c=64, arch_ver=arch_ver) + self.encode = nn.Sequential( + nn.Conv2d(3, 16, 3, 2, 1), nn.ConvTranspose2d(16, 4, 4, 2, 1) + ) + if arch_ver in ["4.10"]: + self.block0 = IFBlock(7 + 16, c=192) + self.block1 = IFBlock(8 + 4 + 16, c=128) + self.block2 = IFBlock(8 + 4 + 16, c=96) + self.block3 = IFBlock(8 + 4 + 16, c=64) + self.encode = nn.Sequential( + nn.Conv2d(3, 32, 3, 2, 1), + nn.LeakyReLU(0.2, True), + nn.Conv2d(32, 32, 3, 1, 1), + nn.LeakyReLU(0.2, True), + nn.Conv2d(32, 32, 3, 1, 1), + nn.LeakyReLU(0.2, True), + nn.ConvTranspose2d(32, 8, 4, 2, 1), + ) + + if arch_ver in ["4.0", "4.2", "4.3"]: + self.contextnet = Contextnet(arch_ver=arch_ver) + self.unet = Unet(arch_ver=arch_ver) + self.arch_ver = arch_ver + + def forward( + self, + img0, + img1, + timestep=0.5, + scale_list=[8, 4, 2, 1], + training=True, + fastmode=True, + ensemble=False, + return_flow=False, + ): + img0 = torch.clamp(img0, 0, 1) + img1 = torch.clamp(img1, 0, 1) + + n, c, h, w = img0.shape + ph = ((h - 1) // 64 + 1) * 64 + pw = ((w - 1) // 64 + 1) * 64 + padding = (0, pw - w, 0, ph - h) + img0 = F.pad(img0, padding) + img1 = F.pad(img1, padding) + x = torch.cat((img0, img1), 1) + + if training == False: + channel = x.shape[1] // 2 + img0 = x[:, :channel] + img1 = x[:, channel:] + if not torch.is_tensor(timestep): + timestep = (x[:, :1].clone() * 0 + 1) * timestep + else: + timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3]) + + flow_list = [] + merged = [] + mask_list = [] + + if self.arch_ver in ["4.7", "4.10"]: + f0 = self.encode(img0[:, :3]) + f1 = self.encode(img1[:, :3]) + + warped_img0 = img0 + warped_img1 = img1 + flow = None + mask = None + block = [self.block0, self.block1, self.block2, self.block3] + + for i in range(4): + if flow is None: + # 4.0-4.6 + if self.arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]: + flow, mask = block[i]( + torch.cat((img0[:, :3], img1[:, :3], timestep), 1), + None, + scale=scale_list[i], + ) + if ensemble: + f1, m1 = block[i]( + torch.cat((img1[:, :3], img0[:, :3], 1 - timestep), 1), + None, + scale=scale_list[i], + ) + flow = (flow + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2 + mask = (mask + (-m1)) / 2 + + # 4.7+ + if self.arch_ver in ["4.7", "4.10"]: + flow, mask = block[i]( + torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1), + None, + scale=scale_list[i], + ) + + if ensemble: + f_, m_ = block[i]( + torch.cat( + (img1[:, :3], img0[:, :3], f1, f0, 1 - timestep), 1 + ), + None, + scale=scale_list[i], + ) + flow = (flow + torch.cat((f_[:, 2:4], f_[:, :2]), 1)) / 2 + mask = (mask + (-m_)) / 2 + + else: + # 4.0-4.6 + if self.arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]: + f0, m0 = block[i]( + torch.cat( + (warped_img0[:, :3], warped_img1[:, :3], timestep, mask), 1 + ), + flow, + scale=scale_list[i], + ) + + if self.arch_ver in ["4.0"]: + if ( + i == 1 + and f0[:, :2].abs().max() > 32 + and f0[:, 2:4].abs().max() > 32 + and not training + ): + for k in range(4): + scale_list[k] *= 2 + flow, mask = block[0]( + torch.cat((img0[:, :3], img1[:, :3], timestep), 1), + None, + scale=scale_list[0], + ) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + f0, m0 = block[i]( + torch.cat( + ( + warped_img0[:, :3], + warped_img1[:, :3], + timestep, + mask, + ), + 1, + ), + flow, + scale=scale_list[i], + ) + + # 4.7+ + if self.arch_ver in ["4.7", "4.10"]: + fd, m0 = block[i]( + torch.cat( + ( + warped_img0[:, :3], + warped_img1[:, :3], + warp(f0, flow[:, :2]), + warp(f1, flow[:, 2:4]), + timestep, + mask, + ), + 1, + ), + flow, + scale=scale_list[i], + ) + flow = flow + fd + + # 4.0-4.6 ensemble + if ensemble and self.arch_ver in [ + "4.0", + "4.2", + "4.3", + "4.5", + "4.6", + ]: + f1, m1 = block[i]( + torch.cat( + ( + warped_img1[:, :3], + warped_img0[:, :3], + 1 - timestep, + -mask, + ), + 1, + ), + torch.cat((flow[:, 2:4], flow[:, :2]), 1), + scale=scale_list[i], + ) + f0 = (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2 + m0 = (m0 + (-m1)) / 2 + + # 4.7+ ensemble + if ensemble and self.arch_ver in ["4.7", "4.10"]: + wf0 = warp(f0, flow[:, :2]) + wf1 = warp(f1, flow[:, 2:4]) + + f_, m_ = block[i]( + torch.cat( + ( + warped_img1[:, :3], + warped_img0[:, :3], + wf1, + wf0, + 1 - timestep, + -mask, + ), + 1, + ), + torch.cat((flow[:, 2:4], flow[:, :2]), 1), + scale=scale_list[i], + ) + fd = (fd + torch.cat((f_[:, 2:4], f_[:, :2]), 1)) / 2 + mask = (m0 + (-m_)) / 2 + + if self.arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]: + flow = flow + f0 + mask = mask + m0 + + if not ensemble and self.arch_ver in ["4.7", "4.10"]: + mask = m0 + + mask_list.append(mask) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + merged.append((warped_img0, warped_img1)) + + if self.arch_ver in ["4.0", "4.1", "4.2", "4.3", "4.4", "4.5", "4.6"]: + mask_list[3] = torch.sigmoid(mask_list[3]) + merged[3] = merged[3][0] * mask_list[3] + merged[3][1] * (1 - mask_list[3]) + + if self.arch_ver in ["4.7", "4.10"]: + mask = torch.sigmoid(mask) + merged[3] = warped_img0 * mask + warped_img1 * (1 - mask) + + if not fastmode and self.arch_ver in ["4.0", "4.2", "4.3"]: + c0 = self.contextnet(img0, flow[:, :2]) + c1 = self.contextnet(img1, flow[:, 2:4]) + tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) + res = tmp[:, :3] * 2 - 1 + merged[3] = torch.clamp(merged[3] + res, 0, 1) + return merged[3][:, :, :h, :w] \ No newline at end of file