Add files via upload
This commit is contained in:
Smai-Lee
@@ -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).
|
||||
|
||||
@@ -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`).
|
||||
|
||||
@@ -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.
|
||||
Binary file not shown.
@@ -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)"}
|
||||
@@ -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).
|
||||
+586
@@ -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]
|
||||
Reference in New Issue
Block a user