Merge pull request #126 from s9901/add-rife-4.26

update RIFE_VFI:supportRIFE 4.26 and RIFE 4.17

Author: Fannovel16

vfi_models/rife/__init__.py

@@ -10,6 +10,8 @@
 CKPT_NAME_VER_DICT = {
     "rife47.pth": "4.7",
     "rife49.pth": "4.7",
+    "rife417.pth": "4.17",
+    "rife426.pth": "4.26",
     "sudo_rife4_269.662_testV1_scale1.pth": "4.0",
     # Arch 4.10 doesn't work due to state dict mismatch
     # "rife410.pth": "4.10",
@@ -118,6 +120,10 @@ def vfi(
         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)
@@ -147,7 +153,11 @@ def vfi(
             multipliers = list(map(int, multiplier))
             multipliers += [2] * (n_pairs - len(multipliers))
 
-        scale_list = [8 / scale_factor, 4 / scale_factor, 2 / scale_factor, 1 / scale_factor]
+         #  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.

vfi_models/rife/rife_arch.py

@@ -92,7 +92,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,
@@ -148,7 +148,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,
@@ -197,7 +197,7 @@ def __init__(self, in_planes, c=64, arch_ver="4.0"):
             )
             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),
@@ -212,11 +212,28 @@ def __init__(self, in_planes, c=64, arch_ver="4.0"):
             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
@@ -233,7 +250,7 @@ def forward(self, x, flow=None, scale=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"]:
+        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)
@@ -242,11 +259,19 @@ def forward(self, x, flow=None, scale=1):
                 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
 
@@ -328,6 +353,48 @@ def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1):
 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):
@@ -347,7 +414,25 @@ def __init__(self, arch_ver="4.0"):
             self.encode = nn.Sequential(
                 nn.Conv2d(3, 16, 3, 2, 1), nn.ConvTranspose2d(16, 4, 4, 2, 1)
             )
-        if arch_ver in ["4.10"]:
+        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)
@@ -362,6 +447,16 @@ def __init__(self, arch_ver="4.0"):
                 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)
@@ -372,7 +467,7 @@ def forward(
         img0,
         img1,
         timestep=0.5,
-        scale_list=[8, 4, 2, 1],
+        scale_list=[16, 8, 4, 2, 1],
         training=True,
         fastmode=True,
         ensemble=False,
@@ -402,20 +497,34 @@ def forward(
         merged = []
         mask_list = []
 
-        if self.arch_ver in ["4.7", "4.10"]:
+        # 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
-        block = [self.block0, self.block1, self.block2, self.block3]
+        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(4):
+        for i in range(num_blocks):
             if flow is None:
-                # 4.0-4.6
-                if self.arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]:
+                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,
@@ -431,7 +540,7 @@ def forward(
                         mask = (mask + (-m1)) / 2
 
                 # 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,
@@ -450,8 +559,34 @@ def forward(
                         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
-                if self.arch_ver in ["4.0", "4.2", "4.3", "4.5", "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
@@ -491,7 +626,7 @@ def forward(
                         )
 
                 # 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(
                             (
@@ -534,7 +669,7 @@ def forward(
                     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])
 
@@ -560,7 +695,7 @@ def forward(
                     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
 
             mask_list.append(mask)
@@ -569,18 +704,29 @@ def forward(
             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)
-        return merged[3][:, :, :h, :w]
+            merged_final = torch.clamp(merged_final + res, 0, 1)
+        
+        return merged_final[:, :, :h, :w]