commit 4a7630404acf8eb24ff5f0f2ecbb3909dce09d56 Author: er1cw00 Date: Tue May 19 23:15:50 2026 +0800 feat: first create diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..9cbc517 --- /dev/null +++ b/__init__.py @@ -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'] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..073d37f --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +torch +torchvision +numpy +einops +urllib diff --git a/rife_model.py b/rife_model.py new file mode 100644 index 0000000..5239e33 --- /dev/null +++ b/rife_model.py @@ -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] diff --git a/rife_node.py b/rife_node.py new file mode 100644 index 0000000..5b8151f --- /dev/null +++ b/rife_node.py @@ -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),)