feat: first create
This commit is contained in:
er1cw00
+11
@@ -0,0 +1,11 @@
|
||||
from .rife_node import RIFE_VFI
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"RIFE_VFI": RIFE_VFI,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"RIFE_VFI": "RIFE_VFI",
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -0,0 +1,5 @@
|
||||
torch
|
||||
torchvision
|
||||
numpy
|
||||
einops
|
||||
urllib
|
||||
+732
@@ -0,0 +1,732 @@
|
||||
"""
|
||||
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 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=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", "4.17", "4.26"]:
|
||||
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", "4.17", "4.26"]:
|
||||
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", "4.17"]:
|
||||
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", "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
|
||||
)
|
||||
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", "4.17", "4.26"]:
|
||||
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", "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
|
||||
|
||||
|
||||
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 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):
|
||||
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", "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()
|
||||
|
||||
if arch_ver in ["old_4.10_delete_me"]:
|
||||
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),
|
||||
)
|
||||
|
||||
# add 4.26 support
|
||||
if arch_ver in ["4.26"]:
|
||||
self.block0 = IFBlock(7+8, c=192, arch_ver=arch_ver)
|
||||
self.block1 = IFBlock(8+4+8+8, c=128, arch_ver=arch_ver)
|
||||
self.block2 = IFBlock(8+4+8+8, c=96, arch_ver=arch_ver)
|
||||
self.block3 = IFBlock(8+4+8+8, c=64, arch_ver=arch_ver)
|
||||
self.block4 = IFBlock(8+4+8+8, c=32, arch_ver=arch_ver)
|
||||
self.encode = Head()
|
||||
|
||||
|
||||
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=[16, 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 = []
|
||||
|
||||
# add 4.26 support
|
||||
if self.arch_ver in ["4.7", "4.10", "4.17", "4.26"]:
|
||||
f0 = self.encode(img0[:, :3])
|
||||
f1 = self.encode(img1[:, :3])
|
||||
|
||||
warped_img0 = img0
|
||||
warped_img1 = img1
|
||||
flow = None
|
||||
mask = None
|
||||
feat = None
|
||||
|
||||
# add 4.26 support
|
||||
if self.arch_ver in ["4.26"]:
|
||||
block = [self.block0, self.block1, self.block2, self.block3, self.block4]
|
||||
num_blocks = 5
|
||||
else:
|
||||
block = [self.block0, self.block1, self.block2, self.block3]
|
||||
num_blocks = 4
|
||||
|
||||
for i in range(num_blocks):
|
||||
if flow is None:
|
||||
if self.arch_ver in ["4.26"]:
|
||||
flow, mask, feat = block[i](
|
||||
torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1),
|
||||
None,
|
||||
scale=scale_list[i],
|
||||
)
|
||||
elif 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", "4.17"]:
|
||||
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:
|
||||
# add 4.26 support
|
||||
if self.arch_ver in ["4.26"]:
|
||||
wf0 = warp(f0, flow[:, :2])
|
||||
wf1 = warp(f1, flow[:, 2:4])
|
||||
|
||||
# add 4.26 support
|
||||
input_tensor = torch.cat(
|
||||
(
|
||||
warped_img0[:, :3],
|
||||
warped_img1[:, :3],
|
||||
wf0,
|
||||
wf1,
|
||||
timestep,
|
||||
mask,
|
||||
feat,
|
||||
),
|
||||
1,
|
||||
)
|
||||
|
||||
fd, m0, feat = block[i](
|
||||
input_tensor,
|
||||
flow,
|
||||
scale=scale_list[i],
|
||||
)
|
||||
flow = flow + fd
|
||||
mask = m0
|
||||
# 4.0-4.6
|
||||
elif 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", "4.17"]:
|
||||
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", "4.17"]:
|
||||
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", "4.17"]:
|
||||
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))
|
||||
|
||||
# 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_final = torch.clamp(merged_final + res, 0, 1)
|
||||
|
||||
return merged_final[:, :, :h, :w]
|
||||
+290
@@ -0,0 +1,290 @@
|
||||
import gc
|
||||
import os
|
||||
import torch
|
||||
import pathlib
|
||||
import einops
|
||||
import folder_paths
|
||||
import typing
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
#from vfi_utils import load_file_from_github_release
|
||||
from comfy.model_management import get_torch_device, soft_empty_cache
|
||||
from packaging import version
|
||||
from .rife_model import IFNet
|
||||
|
||||
MODEL_TYPE = pathlib.Path(__file__).parent.name
|
||||
|
||||
CKPT_NAME_VER_DICT = {
|
||||
"rife47.pth": "4.7",
|
||||
"rife49.pth": "4.7",
|
||||
"rife417.pth": "4.17",
|
||||
}
|
||||
|
||||
CKPT_NAMES = list(CKPT_NAME_VER_DICT.keys())
|
||||
|
||||
DTYPE_OPTIONS = ["float32", "float16", "bfloat16"]
|
||||
DTYPE_MAP = {
|
||||
"float32": torch.float32,
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
}
|
||||
|
||||
# Module-level model cache: avoids reloading weights on every node execution.
|
||||
# Key: (ckpt_name, dtype, torch_compile) — invalidated if any of these change.
|
||||
_model_cache: typing.Dict[typing.Tuple, torch.nn.Module] = {}
|
||||
|
||||
|
||||
def download_model(ckpt_name: str, model_path: str) -> bool:
|
||||
"""Download model from GitHub releases if it doesn't exist locally."""
|
||||
if os.path.exists(model_path):
|
||||
return True
|
||||
|
||||
if ckpt_name not in CKPT_NAME_VER_DICT:
|
||||
print(f"Comfy-VFI: Unknown model '{ckpt_name}', cannot download.")
|
||||
return False
|
||||
|
||||
url = f"https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/releases/download/models/{ckpt_name}"
|
||||
os.makedirs(os.path.dirname(model_path), exist_ok=True)
|
||||
|
||||
try:
|
||||
print(f"Comfy-VFI: Downloading {ckpt_name} from {url}...")
|
||||
urllib.request.urlretrieve(url, model_path)
|
||||
print(f"Comfy-VFI: Successfully downloaded {ckpt_name}")
|
||||
return True
|
||||
except urllib.error.URLError as e:
|
||||
print(f"Comfy-VFI: Failed to download {ckpt_name}: {e}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"Comfy-VFI: Unexpected error downloading {ckpt_name}: {e}")
|
||||
return False
|
||||
|
||||
|
||||
class InterpolationStateList():
|
||||
|
||||
def __init__(self, frame_indices: typing.List[int], is_skip_list: bool):
|
||||
self.frame_indices = frame_indices
|
||||
self.is_skip_list = is_skip_list
|
||||
|
||||
def is_frame_skipped(self, frame_index):
|
||||
is_frame_in_list = frame_index in self.frame_indices
|
||||
return self.is_skip_list and is_frame_in_list or not self.is_skip_list and not is_frame_in_list
|
||||
|
||||
|
||||
def preprocess_frames(frames):
|
||||
return einops.rearrange(frames[..., :3], "n h w c -> n c h w")
|
||||
|
||||
def postprocess_frames(frames):
|
||||
return einops.rearrange(frames, "n c h w -> n h w c")[..., :3].cpu()
|
||||
|
||||
|
||||
class RIFE_VFI:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (CKPT_NAMES,),
|
||||
"frames": ("IMAGE", ),
|
||||
"clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}),
|
||||
"multiplier": ("INT", {"default": 2, "min": 1}),
|
||||
"fast_mode": ("BOOLEAN", {"default": True}),
|
||||
"ensemble": ("BOOLEAN", {"default": True}),
|
||||
"scale_factor": ([0.25, 0.5, 1.0, 2.0, 4.0], {"default": 1.0}),
|
||||
# Opt 1: Precision — float32 is safe, float16 is ~2x faster/half VRAM,
|
||||
# bfloat16 is safer than float16 on Ampere+ GPUs (RTX 30xx+)
|
||||
"dtype": (DTYPE_OPTIONS, {"default": "float32"}),
|
||||
# Opt 2: torch.compile() — JIT-compiles the model graph for 10-30% speedup.
|
||||
# First inference will be slow (compilation), subsequent ones faster.
|
||||
# Requires PyTorch 2.0+. Disable if you hit errors on older builds.
|
||||
"torch_compile": ("BOOLEAN", {"default": False,
|
||||
"tooltip": "Compile the model with torch.compile() for 10-30% faster inference "
|
||||
"after the first (warm-up) run. Requires PyTorch 2.0+."}),
|
||||
# Opt 7: Batch tasks — process multiple (pair, timestep) tasks per GPU call.
|
||||
# Each task is one interpolated frame. Batching amortises kernel launch
|
||||
# overhead; higher values improve throughput but use more VRAM.
|
||||
# 1 = one interpolated frame per call (safest).
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64,
|
||||
"tooltip": "Number of interpolation tasks per GPU call. "
|
||||
"Higher values improve throughput but use more VRAM. "
|
||||
"Set to 1 for the most conservative behaviour."}),
|
||||
},
|
||||
"optional": {
|
||||
"optional_interpolation_states": ("INTERPOLATION_STATES", )
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", )
|
||||
FUNCTION = "vfi"
|
||||
CATEGORY = "ComfyUI-Frame-Interpolation/VFI"
|
||||
|
||||
def _get_model_path(self, model_name: str) -> str:
|
||||
"""获取模型文件的完整路径"""
|
||||
rife_models_folder = os.path.join(folder_paths.models_dir, "rife")
|
||||
return os.path.join(rife_models_folder, model_name)
|
||||
|
||||
def vfi(
|
||||
self,
|
||||
ckpt_name: typing.AnyStr,
|
||||
frames: torch.Tensor,
|
||||
clear_cache_after_n_frames: int = 10,
|
||||
multiplier: typing.SupportsInt = 2,
|
||||
fast_mode: bool = False,
|
||||
ensemble: bool = False,
|
||||
scale_factor: float = 1.0,
|
||||
dtype: str = "bfloat16",
|
||||
torch_compile: bool = False,
|
||||
batch_size: int = 1,
|
||||
optional_interpolation_states: InterpolationStateList = None,
|
||||
**kwargs
|
||||
):
|
||||
"""
|
||||
Perform video frame interpolation using a given checkpoint model.
|
||||
|
||||
Args:
|
||||
ckpt_name (str): The name of the checkpoint model to use.
|
||||
frames (torch.Tensor): A tensor containing input video frames.
|
||||
clear_cache_after_n_frames (int, optional): The number of frames to process before clearing CUDA cache
|
||||
to prevent memory overflow. Defaults to 10. Lower numbers are safer but mean more processing time.
|
||||
How high you should set it depends on how many input frames there are, input resolution (after upscaling),
|
||||
how many times you want to multiply them, and how long you're willing to wait for the process to complete.
|
||||
multiplier (int, optional): The multiplier for each input frame. 60 input frames * 2 = 120 output frames. Defaults to 2.
|
||||
dtype (str, optional): Floating-point precision for inference.
|
||||
"float32" — full precision, safest (default).
|
||||
"float16" — half precision, ~2x faster and ~50% less VRAM, may have minor artifacts.
|
||||
"bfloat16" — brain float16, better numeric range than float16, requires Ampere+ GPU (RTX 30xx+).
|
||||
torch_compile (bool, optional): Compile the model with torch.compile() for 10-30% speedup.
|
||||
The first inference call will be slow (compilation warm-up); subsequent calls are faster.
|
||||
Requires PyTorch 2.0+. Defaults to False.
|
||||
batch_size (int, optional): Number of interpolation tasks per GPU call. Each task produces one
|
||||
intermediate frame. Higher values improve throughput but use more VRAM. Defaults to 1.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing the output interpolated frames.
|
||||
"""
|
||||
|
||||
model_path = self._get_model_path(ckpt_name)
|
||||
arch_ver = CKPT_NAME_VER_DICT[ckpt_name]
|
||||
|
||||
# Download model if it doesn't exist
|
||||
if not download_model(ckpt_name, model_path):
|
||||
raise RuntimeError(f"Model {ckpt_name} not found and could not be downloaded.")
|
||||
|
||||
torch_dtype = DTYPE_MAP[dtype]
|
||||
device = get_torch_device()
|
||||
|
||||
# 4.26 disable ensemble
|
||||
if arch_ver == "4.26":
|
||||
ensemble = False # 4.26 版本不支持 ensemble
|
||||
|
||||
# Cache the model by (ckpt_name, dtype, torch_compile) so repeated node
|
||||
# executions skip the load_state_dict + device transfer entirely.
|
||||
cache_key = (ckpt_name, dtype, torch_compile)
|
||||
if cache_key not in _model_cache:
|
||||
interpolation_model = IFNet(arch_ver=arch_ver)
|
||||
interpolation_model.load_state_dict(torch.load(model_path, map_location=device, weights_only=False))
|
||||
if torch_dtype != torch.float32:
|
||||
interpolation_model = interpolation_model.to(torch_dtype)
|
||||
interpolation_model.eval().to(device)
|
||||
# Opt 2: torch.compile() — JIT-compiles the model for 10-30% faster inference.
|
||||
# The first call triggers compilation (slow); all subsequent calls use the cached graph.
|
||||
if torch_compile:
|
||||
interpolation_model = torch.compile(interpolation_model)
|
||||
_model_cache[cache_key] = interpolation_model
|
||||
print(f"Comfy-VFI: Loaded and cached model {ckpt_name} ({dtype}{'+ torch.compile' if torch_compile else ''})")
|
||||
else:
|
||||
interpolation_model = _model_cache[cache_key]
|
||||
print(f"Comfy-VFI: Using cached model {ckpt_name} ({dtype}{'+ torch.compile' if torch_compile else ''})")
|
||||
|
||||
frames = preprocess_frames(frames)
|
||||
|
||||
# Normalise multiplier to a per-pair list
|
||||
n_pairs = len(frames) - 1
|
||||
if isinstance(multiplier, int):
|
||||
multipliers = [int(multiplier)] * n_pairs
|
||||
else:
|
||||
multipliers = list(map(int, multiplier))
|
||||
multipliers += [2] * (n_pairs - len(multipliers))
|
||||
|
||||
# 4.26 scale_list
|
||||
if arch_ver == "4.26":
|
||||
scale_list = [16 / scale_factor, 8 / scale_factor, 4 / scale_factor, 2 / scale_factor, 1 / scale_factor]
|
||||
else:
|
||||
scale_list = [8 / scale_factor, 4 / scale_factor, 2 / scale_factor, 1 / scale_factor]
|
||||
|
||||
# Build a flat list of all (pair_idx, timestep) tasks, skipping excluded pairs.
|
||||
# Each task produces exactly one intermediate frame.
|
||||
tasks: typing.List[typing.Tuple[int, float]] = []
|
||||
tasks_remaining_per_pair: typing.Dict[int, int] = {}
|
||||
for pair_idx in range(n_pairs):
|
||||
if optional_interpolation_states is not None and optional_interpolation_states.is_frame_skipped(pair_idx):
|
||||
tasks_remaining_per_pair[pair_idx] = 0
|
||||
continue
|
||||
m = multipliers[pair_idx]
|
||||
n_steps = max(m - 1, 0)
|
||||
tasks_remaining_per_pair[pair_idx] = n_steps
|
||||
for step in range(1, m):
|
||||
tasks.append((pair_idx, step / m))
|
||||
|
||||
# Storage for intermediate frames, keyed by pair index
|
||||
results: typing.Dict[int, typing.List[torch.Tensor]] = {i: [] for i in range(n_pairs)}
|
||||
|
||||
frames_processed_since_cache_clear = 0
|
||||
pos = 0
|
||||
|
||||
# Opt 5: inference_mode disables autograd tracking entirely — lower overhead than no_grad().
|
||||
# Opt 7: process batch_size tasks per GPU call. IFNet supports batched tensor timesteps,
|
||||
# so multiple (pair, timestep) combinations can be inferred in a single forward pass.
|
||||
with torch.inference_mode():
|
||||
while pos < len(tasks):
|
||||
batch_tasks = tasks[pos : pos + batch_size]
|
||||
|
||||
frame0_list, frame1_list, timestep_list = [], [], []
|
||||
for pair_idx, dt in batch_tasks:
|
||||
frame0_list.append(frames[pair_idx : pair_idx + 1])
|
||||
frame1_list.append(frames[pair_idx + 1 : pair_idx + 2])
|
||||
timestep_list.append(dt)
|
||||
|
||||
frame0_batch = torch.cat(frame0_list, dim=0).to(device, dtype=torch_dtype)
|
||||
frame1_batch = torch.cat(frame1_list, dim=0).to(device, dtype=torch_dtype)
|
||||
# Batched timestep tensor: shape [B, 1, 1, 1] — IFNet expands it internally
|
||||
timestep_tensor = torch.tensor(timestep_list, dtype=torch_dtype, device=device).view(-1, 1, 1, 1)
|
||||
|
||||
middle_frames = interpolation_model(
|
||||
frame0_batch,
|
||||
frame1_batch,
|
||||
timestep_tensor,
|
||||
scale_list,
|
||||
fast_mode,
|
||||
ensemble,
|
||||
).clamp(0, 1).detach().cpu()
|
||||
|
||||
for idx, (pair_idx, _) in enumerate(batch_tasks):
|
||||
results[pair_idx].append(middle_frames[idx : idx + 1].to(dtype=torch_dtype))
|
||||
tasks_remaining_per_pair[pair_idx] -= 1
|
||||
# Clear cache after finishing each pair's worth of tasks
|
||||
if tasks_remaining_per_pair[pair_idx] == 0:
|
||||
frames_processed_since_cache_clear += 1
|
||||
if frames_processed_since_cache_clear >= clear_cache_after_n_frames:
|
||||
print("Comfy-VFI: Clearing cache...", end=' ')
|
||||
soft_empty_cache()
|
||||
gc.collect()
|
||||
frames_processed_since_cache_clear = 0
|
||||
print("Done cache clearing")
|
||||
|
||||
pos += len(batch_tasks)
|
||||
|
||||
# Assemble output: each original frame followed by its interpolated frames
|
||||
output_frames: typing.List[torch.Tensor] = []
|
||||
for pair_idx in range(n_pairs):
|
||||
output_frames.append(frames[pair_idx : pair_idx + 1].to(dtype=torch_dtype))
|
||||
for mid in results[pair_idx]:
|
||||
output_frames.append(mid)
|
||||
output_frames.append(frames[-1:].to(dtype=torch_dtype))
|
||||
|
||||
print("Comfy-VFI: Final clearing cache...", end=' ')
|
||||
soft_empty_cache()
|
||||
print("Done cache clearing")
|
||||
print(f"Comfy-VFI done! {sum(len(v) for v in results.values()) + len(frames)} frames generated")
|
||||
|
||||
# Always return float32 — numpy and all downstream ComfyUI nodes require it
|
||||
out_tensor = torch.cat(output_frames, dim=0).to(torch.float32)
|
||||
return (postprocess_frames(out_tensor),)
|
||||
Reference in New Issue
Block a user