Add RIFE v4.17, v4.26 and sudo_rife4 models from Fannovel16
Update rife_arch.py with support for new architectures: - v4.17: Head_417 encoder, IFBlock with 16 extra channels - v4.26: 5 IFBlocks (block4 extra), lastconv 4*13 channels, feat_out return, no ensemble, 5 scale levels - Head/Head_417 classes for v4.17/v4.26 encoders - warp() fix: MPS padding_mode, dtype mismatch, get_torch_device() Update export_onnx.py: - Add rife417.pth and rife426.pth to CKPT_NAME_VER_DICT - Add CKPT_FALLBACK_URLS with HuggingFace mirrors for all models - Add export calls for rife49, rife417, rife426, sudo_rife4 Update load_rife_config.json with 3 new ONNX model options. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
huchukato
+43
-6
@@ -23,6 +23,8 @@ CKPT_NAME_VER_DICT = {
|
||||
"rife47.pth": "4.7",
|
||||
"rife48.pth": "4.7",
|
||||
"rife49.pth": "4.7",
|
||||
"rife417.pth": "4.17",
|
||||
"rife426.pth": "4.26",
|
||||
"sudo_rife4_269.662_testV1_scale1.pth": "4.0"
|
||||
}
|
||||
BASE_MODEL_DOWNLOAD_URLS = [
|
||||
@@ -30,6 +32,34 @@ BASE_MODEL_DOWNLOAD_URLS = [
|
||||
"https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/releases/download/models/",
|
||||
"https://github.com/dajes/frame-interpolation-pytorch/releases/download/v1.0.0/"
|
||||
]
|
||||
# Per-file fallback URLs for models no longer hosted at the base URLs above.
|
||||
CKPT_FALLBACK_URLS = {
|
||||
"rife417.pth": [
|
||||
"https://huggingface.co/marduk191/rife/resolve/main/rife417.pth",
|
||||
"https://huggingface.co/MachineDelusions/RIFE/resolve/main/rife417.pth",
|
||||
],
|
||||
"rife426.pth": [
|
||||
"https://huggingface.co/marduk191/rife/resolve/main/rife426.pth",
|
||||
"https://huggingface.co/MachineDelusions/RIFE/resolve/main/rife426.pth",
|
||||
],
|
||||
"rife47.pth": [
|
||||
"https://huggingface.co/marduk191/rife/resolve/main/rife47.pth",
|
||||
"https://huggingface.co/wavespeed/misc/resolve/main/rife/rife47.pth",
|
||||
"https://huggingface.co/MachineDelusions/RIFE/resolve/main/rife47.pth",
|
||||
"https://huggingface.co/jasonot/mycomfyui/resolve/main/rife47.pth",
|
||||
],
|
||||
"rife49.pth": [
|
||||
"https://huggingface.co/marduk191/rife/resolve/main/rife49.pth",
|
||||
"https://huggingface.co/hfmaster/models-moved/resolve/main/rife/rife49.pth",
|
||||
"https://huggingface.co/MachineDelusions/RIFE/resolve/main/rife49.pth",
|
||||
"https://huggingface.co/Isi99999/Frame_Interpolation_Models/resolve/main/rife49.pth",
|
||||
],
|
||||
"sudo_rife4_269.662_testV1_scale1.pth": [
|
||||
"https://huggingface.co/marduk191/rife/resolve/main/sudo_rife4_269.662_testV1_scale1.pth",
|
||||
"https://huggingface.co/uwg/upscaler/resolve/main/ESRGAN/sudo_rife4_269.662_testV1_scale1.pth",
|
||||
"https://huggingface.co/licyk/sd-upscaler-models/resolve/main/ESRGAN/sudo_rife4_269.662_testV1_scale1.pth",
|
||||
],
|
||||
}
|
||||
TORCH_DEVICE = "cuda:0"
|
||||
|
||||
def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
|
||||
@@ -66,17 +96,20 @@ def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
|
||||
def load_file_from_github_release(model_type, ckpt_name):
|
||||
os.makedirs("models", exist_ok=True)
|
||||
error_strs = []
|
||||
for i, base_model_download_url in enumerate(BASE_MODEL_DOWNLOAD_URLS):
|
||||
all_urls = [base + ckpt_name for base in BASE_MODEL_DOWNLOAD_URLS]
|
||||
all_urls += CKPT_FALLBACK_URLS.get(ckpt_name, [])
|
||||
|
||||
for i, url in enumerate(all_urls):
|
||||
try:
|
||||
return load_file_from_url(base_model_download_url + ckpt_name, "models")
|
||||
return load_file_from_url(url, "models")
|
||||
except Exception:
|
||||
traceback_str = traceback.format_exc()
|
||||
if i < len(BASE_MODEL_DOWNLOAD_URLS) - 1:
|
||||
if i < len(all_urls) - 1:
|
||||
print("Failed! Trying another endpoint.")
|
||||
error_strs.append(f"Error when downloading from: {base_model_download_url + ckpt_name}\n\n{traceback_str}")
|
||||
error_strs.append(f"Error when downloading from: {url}\n\n{traceback_str}")
|
||||
|
||||
error_str = '\n\n'.join(error_strs)
|
||||
raise Exception(f"Tried all GitHub base urls to download {ckpt_name} but no suceess. Below is the error log:\n\n{error_str}")
|
||||
raise Exception(f"Tried all urls to download {ckpt_name} but no success. Below is the error log:\n\n{error_str}")
|
||||
|
||||
def export_onnx(ckpt_name, ensemble, scale_factor):
|
||||
print(f"PyTorch version: {torch.__version__}")
|
||||
@@ -140,4 +173,8 @@ def export_onnx(ckpt_name, ensemble, scale_factor):
|
||||
print("=> sim ONNX Model check done!")
|
||||
|
||||
|
||||
export_onnx(ckpt_name="rife47.pth", ensemble=True, scale_factor=1)
|
||||
export_onnx(ckpt_name="rife47.pth", ensemble=True, scale_factor=1)
|
||||
export_onnx(ckpt_name="rife49.pth", ensemble=True, scale_factor=1)
|
||||
export_onnx(ckpt_name="rife417.pth", ensemble=True, scale_factor=1)
|
||||
export_onnx(ckpt_name="rife426.pth", ensemble=False, scale_factor=1)
|
||||
export_onnx(ckpt_name="sudo_rife4_269.662_testV1_scale1.pth", ensemble=True, scale_factor=1)
|
||||
@@ -3,10 +3,13 @@
|
||||
"options": [
|
||||
"rife47_ensemble_True_scale_1_sim",
|
||||
"rife48_ensemble_True_scale_1_sim",
|
||||
"rife49_ensemble_True_scale_1_sim"
|
||||
"rife49_ensemble_True_scale_1_sim",
|
||||
"rife417_ensemble_True_scale_1_sim",
|
||||
"rife426_ensemble_False_scale_1_sim",
|
||||
"sudo_rife4_269.662_testV1_scale1_ensemble_True_scale_1_sim"
|
||||
],
|
||||
"default": "rife49_ensemble_True_scale_1_sim",
|
||||
"tooltip": "RIFE models for video frame interpolation. These models have been tested with tensorrt. Loaded from config."
|
||||
"tooltip": "RIFE models for video frame interpolation. rife49 (v4.7) is the stable default. rife417 (v4.17) and rife426 (v4.26) are newer with improved quality. sudo_rife4 is a community-trained v4.0 variant."
|
||||
},
|
||||
"precision": {
|
||||
"options": ["fp16", "fp32"],
|
||||
|
||||
+169
-41
@@ -13,8 +13,9 @@ 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 = "cuda:0"
|
||||
device = get_torch_device()
|
||||
backwarp_tenGrid = {}
|
||||
|
||||
|
||||
@@ -54,14 +55,19 @@ def warp(tenInput, tenFlow):
|
||||
|
||||
g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1)
|
||||
|
||||
if tenInput.type() == "torch.cuda.HalfTensor":
|
||||
g = g.half()
|
||||
if g.dtype != tenInput.dtype:
|
||||
g = g.to(tenInput.dtype)
|
||||
|
||||
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="border",
|
||||
padding_mode=padding_mode,
|
||||
align_corners=True,
|
||||
)
|
||||
|
||||
@@ -88,7 +94,7 @@ def conv(
|
||||
),
|
||||
nn.PReLU(out_planes),
|
||||
)
|
||||
if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]:
|
||||
if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10", "4.17", "4.26"]:
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_planes,
|
||||
@@ -144,7 +150,7 @@ def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1, arch_ver="
|
||||
),
|
||||
nn.PReLU(out_planes),
|
||||
)
|
||||
if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]:
|
||||
if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10", "4.17", "4.26"]:
|
||||
return nn.Sequential(
|
||||
torch.nn.ConvTranspose2d(
|
||||
in_channels=in_planes,
|
||||
@@ -193,7 +199,7 @@ class IFBlock(nn.Module):
|
||||
)
|
||||
self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)
|
||||
|
||||
if arch_ver in ["4.5", "4.6", "4.7", "4.10"]:
|
||||
if arch_ver in ["4.5", "4.6", "4.7", "4.10", "4.17"]:
|
||||
self.convblock = nn.Sequential(
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
@@ -208,11 +214,28 @@ class IFBlock(nn.Module):
|
||||
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"]:
|
||||
if arch_ver in ["4.6", "4.7", "4.10", "4.17"]:
|
||||
self.lastconv = nn.Sequential(
|
||||
nn.ConvTranspose2d(c, 4 * 6, 4, 2, 1), nn.PixelShuffle(2)
|
||||
)
|
||||
|
||||
# add 4.26 support
|
||||
if arch_ver in ["4.26"]:
|
||||
self.convblock = nn.Sequential(
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
)
|
||||
self.lastconv = nn.Sequential(
|
||||
nn.ConvTranspose2d(c, 4*13, 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
|
||||
@@ -229,7 +252,7 @@ class IFBlock(nn.Module):
|
||||
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"]:
|
||||
if self.arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10", "4.17", "4.26"]:
|
||||
feat = self.convblock(feat)
|
||||
|
||||
tmp = self.lastconv(feat)
|
||||
@@ -238,11 +261,19 @@ class IFBlock(nn.Module):
|
||||
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"]:
|
||||
if self.arch_ver in ["4.5", "4.6", "4.7", "4.10", "4.17"]:
|
||||
tmp = F.interpolate(
|
||||
tmp, scale_factor=scale, mode="bilinear", align_corners=False
|
||||
)
|
||||
flow = tmp[:, :4] * scale
|
||||
if self.arch_ver in ["4.26"]:
|
||||
tmp = F.interpolate(
|
||||
tmp, scale_factor=scale, mode="bilinear", align_corners=False
|
||||
)
|
||||
flow = tmp[:, :4] * scale
|
||||
mask = tmp[:, 4:5]
|
||||
feat_out = tmp[:, 5:]
|
||||
return flow, mask, feat_out
|
||||
mask = tmp[:, 4:5]
|
||||
return flow, mask
|
||||
|
||||
@@ -314,7 +345,7 @@ class Unet(nn.Module):
|
||||
|
||||
|
||||
"""
|
||||
currently supports 4.0-4.12
|
||||
currently supports 4.0-4.12, 4.17, 4.26
|
||||
|
||||
4.0: 4.0, 4.1
|
||||
4.2: 4.2
|
||||
@@ -323,7 +354,51 @@ currently supports 4.0-4.12
|
||||
4.6: 4.6
|
||||
4.7: 4.7, 4.8, 4.9
|
||||
4.10: 4.10 4.11 4.12
|
||||
4.17: 4.17
|
||||
4.26: 4.26
|
||||
"""
|
||||
class Head_417(nn.Module):
|
||||
def __init__(self):
|
||||
super(Head_417, self).__init__()
|
||||
self.cnn0 = nn.Conv2d(3, 32, 3, 2, 1)
|
||||
self.cnn1 = nn.Conv2d(32, 32, 3, 1, 1)
|
||||
self.cnn2 = nn.Conv2d(32, 32, 3, 1, 1)
|
||||
self.cnn3 = nn.ConvTranspose2d(32, 8, 4, 2, 1)
|
||||
self.relu = nn.LeakyReLU(0.2, True)
|
||||
|
||||
def forward(self, x, feat=False):
|
||||
x0 = self.cnn0(x)
|
||||
x = self.relu(x0)
|
||||
x1 = self.cnn1(x)
|
||||
x = self.relu(x1)
|
||||
x2 = self.cnn2(x)
|
||||
x = self.relu(x2)
|
||||
x3 = self.cnn3(x)
|
||||
if feat:
|
||||
return [x0, x1, x2, x3]
|
||||
return x3
|
||||
|
||||
|
||||
class Head(nn.Module):
|
||||
def __init__(self):
|
||||
super(Head, self).__init__()
|
||||
self.cnn0 = nn.Conv2d(3, 16, 3, 2, 1)
|
||||
self.cnn1 = nn.Conv2d(16, 16, 3, 1, 1)
|
||||
self.cnn2 = nn.Conv2d(16, 16, 3, 1, 1)
|
||||
self.cnn3 = nn.ConvTranspose2d(16, 4, 4, 2, 1)
|
||||
self.relu = nn.LeakyReLU(0.2, True)
|
||||
|
||||
def forward(self, x, feat=False):
|
||||
x0 = self.cnn0(x)
|
||||
x = self.relu(x0)
|
||||
x1 = self.cnn1(x)
|
||||
x = self.relu(x1)
|
||||
x2 = self.cnn2(x)
|
||||
x = self.relu(x2)
|
||||
x3 = self.cnn3(x)
|
||||
if feat:
|
||||
return [x0, x1, x2, x3]
|
||||
return x3
|
||||
|
||||
|
||||
class IFNet(nn.Module):
|
||||
@@ -343,20 +418,31 @@ class IFNet(nn.Module):
|
||||
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.10", "4.17"]:
|
||||
self.block0 = IFBlock(7 + 16, c=192, arch_ver=arch_ver)
|
||||
self.block1 = IFBlock(8 + 4 + 16, c=128, arch_ver=arch_ver)
|
||||
self.block2 = IFBlock(8 + 4 + 16, c=96, arch_ver=arch_ver)
|
||||
self.block3 = IFBlock(8 + 4 + 16, c=64, arch_ver=arch_ver)
|
||||
if arch_ver == "4.10":
|
||||
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),
|
||||
)
|
||||
else:
|
||||
self.encode = Head_417()
|
||||
# add 4.26 support
|
||||
if arch_ver in ["4.26"]:
|
||||
self.block0 = IFBlock(7+16, c=192, arch_ver=arch_ver)
|
||||
self.block1 = IFBlock(8+4+16+8, c=128, arch_ver=arch_ver)
|
||||
self.block2 = IFBlock(8+4+16+8, c=96, arch_ver=arch_ver)
|
||||
self.block3 = IFBlock(8+4+16+8, c=64, arch_ver=arch_ver)
|
||||
self.block4 = IFBlock(8+4+16+8, c=64, arch_ver=arch_ver)
|
||||
self.encode = Head_417()
|
||||
|
||||
if arch_ver in ["4.0", "4.2", "4.3"]:
|
||||
self.contextnet = Contextnet(arch_ver=arch_ver)
|
||||
@@ -406,7 +492,7 @@ class IFNet(nn.Module):
|
||||
merged = []
|
||||
mask_list = []
|
||||
|
||||
if self.arch_ver in ["4.7", "4.10"]:
|
||||
if self.arch_ver in ["4.7", "4.10", "4.17", "4.26"]:
|
||||
f0 = self.encode(img0[:, :3])
|
||||
f1 = self.encode(img1[:, :3])
|
||||
|
||||
@@ -416,11 +502,13 @@ class IFNet(nn.Module):
|
||||
flow = None
|
||||
mask = None
|
||||
block = [self.block0, self.block1, self.block2, self.block3]
|
||||
if self.arch_ver in ["4.26"]:
|
||||
block.append(self.block4)
|
||||
|
||||
for i in range(4):
|
||||
for i in range(len(block)):
|
||||
if flow is None:
|
||||
# 4.7+
|
||||
if self.arch_ver in ["4.7", "4.10"]:
|
||||
if self.arch_ver in ["4.7", "4.10", "4.17"]:
|
||||
flow, mask = block[i](
|
||||
torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1),
|
||||
None,
|
||||
@@ -438,9 +526,17 @@ class IFNet(nn.Module):
|
||||
flow = (flow + torch.cat((f_[:, 2:4], f_[:, :2]), 1)) / 2
|
||||
mask = (mask + (-m_)) / 2
|
||||
|
||||
# 4.26
|
||||
if self.arch_ver in ["4.26"]:
|
||||
flow, mask, feat_out = block[i](
|
||||
torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1),
|
||||
None,
|
||||
scale=scale_list[i],
|
||||
)
|
||||
|
||||
else:
|
||||
# 4.7+
|
||||
if self.arch_ver in ["4.7", "4.10"]:
|
||||
if self.arch_ver in ["4.7", "4.10", "4.17"]:
|
||||
fd, m0 = block[i](
|
||||
torch.cat(
|
||||
(
|
||||
@@ -458,6 +554,26 @@ class IFNet(nn.Module):
|
||||
)
|
||||
flow = flow + fd
|
||||
|
||||
# 4.26
|
||||
if self.arch_ver in ["4.26"]:
|
||||
fd, m0, feat_out = block[i](
|
||||
torch.cat(
|
||||
(
|
||||
warped_img0[:, :3],
|
||||
warped_img1[:, :3],
|
||||
warp(f0, flow[:, :2]),
|
||||
warp(f1, flow[:, 2:4]),
|
||||
timestep,
|
||||
mask,
|
||||
feat_out,
|
||||
),
|
||||
1,
|
||||
),
|
||||
flow,
|
||||
scale=scale_list[i],
|
||||
)
|
||||
flow = flow + fd
|
||||
|
||||
# 4.0-4.6 ensemble
|
||||
if ensemble and self.arch_ver in [
|
||||
"4.0",
|
||||
@@ -483,7 +599,7 @@ class IFNet(nn.Module):
|
||||
m0 = (m0 + (-m1)) / 2
|
||||
|
||||
# 4.7+ ensemble
|
||||
if ensemble and self.arch_ver in ["4.7", "4.10"]:
|
||||
if ensemble and self.arch_ver in ["4.7", "4.10", "4.17"]:
|
||||
wf0 = warp(f0, flow[:, :2])
|
||||
wf1 = warp(f1, flow[:, 2:4])
|
||||
|
||||
@@ -509,7 +625,10 @@ class IFNet(nn.Module):
|
||||
flow = flow + f0
|
||||
mask = mask + m0
|
||||
|
||||
if not ensemble and self.arch_ver in ["4.7", "4.10"]:
|
||||
if not ensemble and self.arch_ver in ["4.7", "4.10", "4.17"]:
|
||||
mask = m0
|
||||
|
||||
if self.arch_ver in ["4.26"]:
|
||||
mask = m0
|
||||
|
||||
mask_list.append(mask)
|
||||
@@ -518,20 +637,29 @@ class IFNet(nn.Module):
|
||||
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)
|
||||
# add 4.26 support
|
||||
if self.arch_ver in ["4.26"]:
|
||||
final_mask = mask_list[-1] if mask_list else mask
|
||||
final_mask = torch.sigmoid(final_mask)
|
||||
merged_final = warped_img0 * final_mask + warped_img1 * (1 - final_mask)
|
||||
else:
|
||||
if self.arch_ver in ["4.0", "4.1", "4.2", "4.3", "4.4", "4.5", "4.6"]:
|
||||
final_idx = min(3, len(mask_list)-1) if len(mask_list) > 0 else 0
|
||||
if final_idx < len(mask_list):
|
||||
mask_list[final_idx] = torch.sigmoid(mask_list[final_idx])
|
||||
merged_final = merged[final_idx][0] * mask_list[final_idx] + merged[final_idx][1] * (1 - mask_list[final_idx])
|
||||
else:
|
||||
mask = torch.sigmoid(mask)
|
||||
merged_final = warped_img0 * mask + warped_img1 * (1 - mask)
|
||||
if self.arch_ver in ["4.7", "4.10", "4.17"]:
|
||||
mask = torch.sigmoid(mask)
|
||||
merged_final = 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)
|
||||
merged_final = torch.clamp(merged_final + res, 0, 1)
|
||||
|
||||
result = merged[3][:, :, :h, :w]
|
||||
return result
|
||||
return merged_final[:, :, :h, :w]
|
||||
|
||||
Reference in New Issue
Block a user