From 2e1c533e3482f85b5fc472baf67846ca521a5acd Mon Sep 17 00:00:00 2001 From: huchukato Date: Wed, 5 Aug 2026 11:00:26 +0200 Subject: [PATCH] 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> --- export_onnx.py | 49 ++++++++-- load_rife_config.json | 7 +- rife_arch.py | 210 +++++++++++++++++++++++++++++++++--------- 3 files changed, 217 insertions(+), 49 deletions(-) diff --git a/export_onnx.py b/export_onnx.py index bc34ea6..87426f3 100644 --- a/export_onnx.py +++ b/export_onnx.py @@ -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) \ No newline at end of file +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) \ No newline at end of file diff --git a/load_rife_config.json b/load_rife_config.json index 4d1e141..fb4d714 100644 --- a/load_rife_config.json +++ b/load_rife_config.json @@ -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"], diff --git a/rife_arch.py b/rife_arch.py index 45a2f10..3e7bc62 100644 --- a/rife_arch.py +++ b/rife_arch.py @@ -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 \ No newline at end of file + return merged_final[:, :, :h, :w]