diff --git a/__init__.py b/__init__.py index b388cb2..21da8a7 100644 --- a/__init__.py +++ b/__init__.py @@ -309,8 +309,8 @@ def load_node_config(config_filename="load_rife_config.json"): default_config = { "model": { - "options": ["rife49_ensemble_True_scale_1_sim"], - "default": "rife49_ensemble_True_scale_1_sim", + "options": ["rife425_ensemble_False_scale_1_sim", "rife49_ensemble_True_scale_1_sim"], + "default": "rife425_ensemble_False_scale_1_sim", "tooltip": "Default model (fallback from code)" }, "precision": { diff --git a/export_onnx.py b/export_onnx.py index 0442e60..d010eab 100644 --- a/export_onnx.py +++ b/export_onnx.py @@ -24,6 +24,7 @@ CKPT_NAME_VER_DICT = { "rife48.pth": "4.7", "rife49.pth": "4.7", "rife417.pth": "4.17", + "rife425.pth": "4.25", "rife426.pth": "4.26", "sudo_rife4_269.662_testV1_scale1.pth": "4.0" } @@ -38,6 +39,9 @@ CKPT_FALLBACK_URLS = { "https://huggingface.co/marduk191/rife/resolve/main/rife417.pth", "https://huggingface.co/MachineDelusions/RIFE/resolve/main/rife417.pth", ], + "rife425.pth": [ + "https://huggingface.co/huchukato/favs/resolve/main/RIFE/rife425.pth", + ], "rife426.pth": [ "https://huggingface.co/marduk191/rife/resolve/main/rife426.pth", "https://huggingface.co/MachineDelusions/RIFE/resolve/main/rife426.pth", @@ -60,7 +64,7 @@ CKPT_FALLBACK_URLS = { "https://huggingface.co/licyk/sd-upscaler-models/resolve/main/ESRGAN/sudo_rife4_269.662_testV1_scale1.pth", ], } -TORCH_DEVICE = "cuda:0" +TORCH_DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu" def load_file_from_url(url, model_dir=None, progress=True, file_name=None): """Load file form http url, will download models if necessary. @@ -188,5 +192,6 @@ def export_onnx(ckpt_name, ensemble, scale_factor): 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="rife425.pth", ensemble=False, 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 fb4d714..9b3e116 100644 --- a/load_rife_config.json +++ b/load_rife_config.json @@ -1,15 +1,11 @@ { "model": { "options": [ - "rife47_ensemble_True_scale_1_sim", - "rife48_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" + "rife425_ensemble_False_scale_1_sim" ], - "default": "rife49_ensemble_True_scale_1_sim", - "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." + "default": "rife425_ensemble_False_scale_1_sim", + "tooltip": "RIFE models for video frame interpolation. rife425 (v4.25) is the recommended default for diffusion-generated video. rife49 (v4.9) is the stable legacy option." }, "precision": { "options": ["fp16", "fp32"], diff --git a/requirements.txt b/requirements.txt index 248531e..e883849 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,7 +8,10 @@ # Base dependencies that work with any CUDA version einops colored -polygraphy +# Pin to the Polygraphy version bundled with TensorRT 10.15.1. +# Newer Polygraphy can break trt.OnnxParser init with: +# pybind11::init(): factory function returned nullptr +polygraphy==0.49.26 cuda-python requests tqdm diff --git a/rife_arch.py b/rife_arch.py index 3e7bc62..dba63b2 100644 --- a/rife_arch.py +++ b/rife_arch.py @@ -94,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", "4.17", "4.26"]: + if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10", "4.17", "4.26", "4.25"]: return nn.Sequential( nn.Conv2d( in_planes, @@ -150,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", "4.17", "4.26"]: + if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10", "4.17", "4.26", "4.25"]: return nn.Sequential( torch.nn.ConvTranspose2d( in_channels=in_planes, @@ -219,8 +219,8 @@ class IFBlock(nn.Module): nn.ConvTranspose2d(c, 4 * 6, 4, 2, 1), nn.PixelShuffle(2) ) - # add 4.26 support - if arch_ver in ["4.26"]: + # add 4.26/4.25 support + if arch_ver in ["4.26", "4.25"]: self.convblock = nn.Sequential( ResConv(c), ResConv(c), @@ -252,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", "4.17", "4.26"]: + if self.arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10", "4.17", "4.26", "4.25"]: feat = self.convblock(feat) tmp = self.lastconv(feat) @@ -266,7 +266,7 @@ class IFBlock(nn.Module): tmp, scale_factor=scale, mode="bilinear", align_corners=False ) flow = tmp[:, :4] * scale - if self.arch_ver in ["4.26"]: + if self.arch_ver in ["4.26", "4.25"]: tmp = F.interpolate( tmp, scale_factor=scale, mode="bilinear", align_corners=False ) @@ -443,6 +443,14 @@ class IFNet(nn.Module): 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() + # add 4.25 support (same structure as 4.26 but smaller encode head and block4) + if arch_ver in ["4.25"]: + 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) @@ -464,6 +472,8 @@ class IFNet(nn.Module): img1 = torch.clamp(img1, 0, 1) scale_list = [8,4,2,1] + if self.arch_ver in ["4.25", "4.26"]: + scale_list = [16,8,4,2,1] training=True fastmode=True @@ -492,7 +502,7 @@ class IFNet(nn.Module): merged = [] mask_list = [] - if self.arch_ver in ["4.7", "4.10", "4.17", "4.26"]: + if self.arch_ver in ["4.7", "4.10", "4.17", "4.26", "4.25"]: f0 = self.encode(img0[:, :3]) f1 = self.encode(img1[:, :3]) @@ -502,7 +512,7 @@ class IFNet(nn.Module): flow = None mask = None block = [self.block0, self.block1, self.block2, self.block3] - if self.arch_ver in ["4.26"]: + if self.arch_ver in ["4.26", "4.25"]: block.append(self.block4) for i in range(len(block)): @@ -526,8 +536,8 @@ 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"]: + # 4.26 / 4.25 + if self.arch_ver in ["4.26", "4.25"]: flow, mask, feat_out = block[i]( torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1), None, @@ -554,8 +564,8 @@ class IFNet(nn.Module): ) flow = flow + fd - # 4.26 - if self.arch_ver in ["4.26"]: + # 4.26 / 4.25 + if self.arch_ver in ["4.26", "4.25"]: fd, m0, feat_out = block[i]( torch.cat( ( @@ -628,7 +638,7 @@ class IFNet(nn.Module): if not ensemble and self.arch_ver in ["4.7", "4.10", "4.17"]: mask = m0 - if self.arch_ver in ["4.26"]: + if self.arch_ver in ["4.26", "4.25"]: mask = m0 mask_list.append(mask) @@ -637,8 +647,8 @@ class IFNet(nn.Module): warped_img1 = warp(img1, flow[:, 2:4]) merged.append((warped_img0, warped_img1)) - # add 4.26 support - if self.arch_ver in ["4.26"]: + # add 4.26/4.25 support + if self.arch_ver in ["4.26", "4.25"]: 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)