added model weights and code

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright  (c)  Megvii  Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,107 @@
----
-license: mit
----
+---
+license: mit
+tags:
+  - video-interpolation
+  - frame-interpolation
+  - optical-flow
+  - pytorch
+  - safetensors
+  - tensorrt
+  - torch-compile
+library_name: pytorch
+---
+
+# RIFE — Real-Time Intermediate Flow Estimation
+
+Safetensors re-host of the RIFE v4 frame interpolation model from
+[ECCV2022-RIFE](https://github.com/hzwer/ECCV2022-RIFE) (MIT license, © Megvii Inc.).
+
+## What is different from the original
+
+| | Original | This repo |
+|---|---|---|
+| **Weight format** | `.pkl` hosted on Google Drive | `.safetensors` hosted on Hugging Face |
+| **TensorRT / `torch.compile`** | Known issues with torchinductor and TensorRT backends | Fixed — model sources are fully compatible |
+
+The architectural code in [interpolation_model.py](interpolation_model.py) is a minimal
+clean-up of the upstream
+[`model/IFNet.py`](https://github.com/hzwer/ECCV2022-RIFE/blob/main/model/IFNet.py)
+with all changes required for `torch.compile` and TensorRT export applied.
+
+## Model description
+
+RIFE (Real-Time Intermediate Flow Estimation) estimates an intermediate video frame
+between two input frames by computing bidirectional optical flow and blending the
+warped frames with a learned mask.
+
+**Input:** a `(B, 6, H, W)` tensor — the first frame in channels `[:3]` and the
+second frame in channels `[3:]`, values in `[0, 1]`.  
+**Output:** a `(B, 3, H, W)` interpolated frame tensor, values in `[0, 1]`.
+
+## Installation
+
+```bash
+git clone https://huggingface.co/tensorforger/RIFE-safetensors
+cd RIFE-safetensors
+pip install -r requirements.txt
+```
+
+## Usage
+
+```python
+import torch
+from safetensors.torch import load_file
+from interpolation_model import IFNet
+
+model = IFNet()
+model.load_state_dict(load_file("flownet.safetensors"))
+model.to("cuda").eval()
+
+# frame0, frame1: (B, 3, H, W) float tensors in [0, 1]
+x = torch.cat([frame0, frame1], dim=1)   # → (B, 6, H, W)
+with torch.no_grad():
+    mid_frame = model(x)                 # → (B, 3, H, W)
+```
+
+### torch.compile (torchinductor / TensorRT)
+
+```python
+model = torch.compile(model, backend="inductor")   # or backend="tensorrt"
+```
+
+### Run the bundled demo
+
+The demo generates two synthetic frames with shifted gray squares and displays the
+interpolated result with `matplotlib`.
+
+```bash
+python demo.py
+```
+
+Output shape printed to stdout: `torch.Size([1, 3, 256, 256])`.  
+A window will open showing **frame 0 · interpolated frame · frame 1**.
+
+## Files
+
+| File | Description |
+|------|-------------|
+| `flownet.safetensors` | Model weights (converted from original `.pkl`) |
+| `interpolation_model.py` | `IFNet` model definition (compile-friendly fork of upstream) |
+| `demo.py` | Minimal runnable example |
+| `requirements.txt` | Python dependencies |
+
+## Citation
+
+```bibtex
+@inproceedings{huang2022rife,
+  title     = {Real-Time Intermediate Flow Estimation for Video Frame Interpolation},
+  author    = {Huang, Zhewei and Zhang, Tianyuan and Heng, Wen and Shi, Boxin and Zhou, Shuchang},
+  booktitle = {Proceedings of the European Conference on Computer Vision (ECCV)},
+  year      = {2022}
+}
+```
+
+## License
+
+MIT — see [LICENSE](LICENSE).  
+Original work © Megvii Inc. This re-host adds no new restrictions.

demo.py

@@ -0,0 +1,28 @@
+import matplotlib.pyplot as plt
+import torch
+from safetensors.torch import load_file
+from interpolation_model import IFNet
+
+interpolation_model = IFNet()
+interpolation_model.load_state_dict(load_file("flownet.safetensors"))
+interpolation_model.to("cuda")
+interpolation_model.eval()
+
+input_frames = torch.zeros([1, 6, 256, 256], device="cuda")
+
+# draw squares
+input_frames[:, :3, 100:200, 100:200] = 0.5
+input_frames[:, 3:, 140:250, 140:240] = 0.5
+
+with torch.no_grad():
+    interpolated_frame = interpolation_model(input_frames)
+
+print(interpolated_frame.shape)  # torch.Size([1, 3, 256, 256])
+
+fig, ax = plt.subplots(1, 3, figsize=(12, 4))
+
+ax[0].imshow(input_frames[0, :3].permute(1, 2, 0).cpu().numpy())
+ax[1].imshow(interpolated_frame[0].permute(1, 2, 0).cpu().numpy())
+ax[2].imshow(input_frames[0, 3:].permute(1, 2, 0).cpu().numpy())
+
+plt.show()

flownet.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b953bee2017bea5e71aae69c36a9caf316071788f9956b7099c015c75ee6af7
+size 12166924

interpolation_model.py

@@ -0,0 +1,184 @@
+# Forked from https://github.com/hzwer/ECCV2022-RIFE/blob/main/model/IFNet.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dtype = torch.float16
+
+
+def conv(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,
+        ),
+        nn.PReLU(out_planes),
+    )
+
+
+def conv_bn(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=False,
+        ),
+        nn.BatchNorm2d(out_planes),
+        nn.PReLU(out_planes),
+    )
+
+
+class IFBlock(nn.Module):
+    def __init__(self, in_planes, c=64):
+        super(IFBlock, self).__init__()
+        self.conv0 = nn.Sequential(
+            conv(in_planes, c // 2, 3, 2, 1),
+            conv(c // 2, c, 3, 2, 1),
+        )
+        self.convblock0 = nn.Sequential(conv(c, c), conv(c, c))
+        self.convblock1 = nn.Sequential(conv(c, c), conv(c, c))
+        self.convblock2 = nn.Sequential(conv(c, c), conv(c, c))
+        self.convblock3 = nn.Sequential(conv(c, c), conv(c, c))
+        self.conv1 = nn.Sequential(
+            nn.ConvTranspose2d(c, c // 2, 4, 2, 1),
+            nn.PReLU(c // 2),
+            nn.ConvTranspose2d(c // 2, 4, 4, 2, 1),
+        )
+        self.conv2 = nn.Sequential(
+            nn.ConvTranspose2d(c, c // 2, 4, 2, 1),
+            nn.PReLU(c // 2),
+            nn.ConvTranspose2d(c // 2, 1, 4, 2, 1),
+        )
+
+    def forward(self, x, flow, scale=1):
+        x = F.interpolate(
+            x,
+            scale_factor=1.0 / scale,
+            mode="bilinear",
+            align_corners=False,
+            recompute_scale_factor=False,
+        )
+        flow = (
+            F.interpolate(
+                flow,
+                scale_factor=1.0 / scale,
+                mode="bilinear",
+                align_corners=False,
+                recompute_scale_factor=False,
+            )
+            * 1.0
+            / scale
+        )
+        feat = self.conv0(torch.cat((x, flow), 1))
+        feat = self.convblock0(feat) + feat
+        feat = self.convblock1(feat) + feat
+        feat = self.convblock2(feat) + feat
+        feat = self.convblock3(feat) + feat
+        flow = self.conv1(feat)
+        mask = self.conv2(feat)
+        flow = (
+            F.interpolate(
+                flow,
+                scale_factor=scale,
+                mode="bilinear",
+                align_corners=False,
+                recompute_scale_factor=False,
+            )
+            * scale
+        )
+        mask = F.interpolate(
+            mask,
+            scale_factor=scale,
+            mode="bilinear",
+            align_corners=False,
+            recompute_scale_factor=False,
+        )
+        return flow, mask
+
+
+class IFNet(nn.Module):
+    def __init__(self):
+        super(IFNet, self).__init__()
+        self.block0 = IFBlock(7 + 4, c=90)
+        self.block1 = IFBlock(7 + 4, c=90)
+        self.block2 = IFBlock(7 + 4, c=90)
+        self.block_tea = IFBlock(10 + 4, c=90)
+
+    def forward(self, x):
+        scale_list = [4, 2, 1]
+
+        channel = x.shape[1] // 2
+        img0 = x[:, :channel]
+        img1 = x[:, channel:]
+        flow_list = []
+        merged = []
+        mask_list = []
+        warped_img0 = img0
+        warped_img1 = img1
+        flow = (x[:, :4]).detach() * 0
+        mask = (x[:, :1]).detach() * 0
+        loss_cons = 0
+        block = [self.block0, self.block1, self.block2]
+        for i in range(3):
+            f0, m0 = block[i](
+                torch.cat((warped_img0[:, :3], warped_img1[:, :3], mask), 1),
+                flow,
+                scale=scale_list[i],
+            )
+            f1, m1 = block[i](
+                torch.cat((warped_img1[:, :3], warped_img0[:, :3], -mask), 1),
+                torch.cat((flow[:, 2:4], flow[:, :2]), 1),
+                scale=scale_list[i],
+            )
+            flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2
+            mask = mask + (m0 + (-m1)) / 2
+            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))
+
+        for i in range(3):
+            mask_list[i] = torch.sigmoid(mask_list[i])
+            merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])
+        return merged[2]
+
+
+def warp(tenInput, tenFlow):
+    tenHorizontal = (
+        torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device, dtype=dtype)
+        .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, dtype=dtype)
+        .view(1, 1, tenFlow.shape[2], 1)
+        .expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3])
+    )
+    backwarp_tenGrid = torch.cat([tenHorizontal, tenVertical], 1).to(device, dtype)
+    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 + tenFlow).permute(0, 2, 3, 1)
+    return torch.nn.functional.grid_sample(
+        input=tenInput,
+        grid=g,
+        mode="bilinear",
+        padding_mode="border",
+        align_corners=True,
+    )

requirements.txt

@@ -0,0 +1,4 @@
+torch>=1.12.0
+safetensors>=0.4.0
+numpy>=1.21.0
+matplotlib>=3.5.0