How ContextNet and UNet Modules Work in the RIFE Architecture

ContextNet extracts multi-scale, flow-warped context features from input frames, while UNet fuses these features with warped images and optical flow to predict the final interpolated frame and blending mask.

The RIFE (Real-time Intermediate Flow Estimation) architecture by hzwer/eccv2022-rife achieves high-quality video frame interpolation through a three-stage pipeline: coarse flow estimation, context-aware feature extraction, and multi-modal fusion. While IFNet handles the initial motion estimation, the ContextNet and UNet modules collaborate to align spatial features and refine the output. Understanding their interaction is essential for adapting the model to custom resolutions or modifying the fusion behavior.

ContextNet – Hierarchical Flow-Aware Feature Extraction

ContextNet (and its lightweight variant Contextnet) generates spatially-aligned feature pyramids that enable the UNet to reason about both fine textures and global scene structure. The module processes each input frame independently, warping features at multiple scales using the intermediate optical flow predicted by IFNet.

Architecture Implementation in RIFE_HDv2

The high-capacity version resides in model/oldmodel/RIFE_HDv2.py and uses a channel factor c = 32:

class ContextNet(nn.Module):
    def __init__(self):
        super(ContextNet, self).__init__()
        self.conv0 = Conv2(3, c)          # 3→c   (c=32)

        self.conv1 = Conv2(c, c)          # c→c

        self.conv2 = Conv2(c, 2*c)        # c→2c

        self.conv3 = Conv2(2*c, 4*c)      # 2c→4c

        self.conv4 = Conv2(4*c, 8*c)      # 4c→8c

    def forward(self, x, flow):
        x = self.conv0(x)
        x = self.conv1(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f1 = warp(x, flow)                # warp at 1/4 resolution

        x = self.conv2(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f2 = warp(x, flow)                # warp at 1/8 resolution

        x = self.conv3(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f3 = warp(x, flow)                # warp at 1/16 resolution

        x = self.conv4(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f4 = warp(x, flow)                # warp at 1/32 resolution

        return [f1, f2, f3, f4]

Each Conv2 block performs stride-2 downsampling while doubling the channel depth, creating a pyramid with resolutions at 1/4, 1/8, 1/16, and 1/32 of the original.

Lightweight Variant in refine.py

The newer implementation in model/refine.py uses the same logic with reduced capacity (c = 16) to improve inference speed:

c = 16
class Contextnet(nn.Module):
    def __init__(self):
        super(Contextnet, self).__init__()
        self.conv1 = Conv2(3, c)
        self.conv2 = Conv2(c, 2*c)
        self.conv3 = Conv2(2*c, 4*c)
        self.conv4 = Conv2(4*c, 8*c)

    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]

Flow-Guided Warping at Multiple Scales

After every downsampling stage, ContextNet applies the differentiable warp function (implemented in model/warplayer.py) to align features with the intermediate frame. The flow tensor is progressively halved using F.interpolate and scaled by 0.5 to maintain correct motion vector magnitudes as spatial resolution decreases. This produces four warped feature maps that preserve semantic context despite large motion between frames.

UNet – Fusion and Refinement Network

The UNet (class Unet in model/refine.py, also referred to as FusionNet in earlier versions) aggregates information from the warped frames, predicted flow, soft masks, and the multi-scale context features generated by ContextNet.

Encoder Path with Context Injection

The encoder consists of four downsampling blocks. The first layer receives a 17-channel tensor concatenating both original frames, their warped versions, the soft mask, and the bidirectional flow:

class Unet(nn.Module):
    def __init__(self):
        super(Unet, self).__init__()
        self.down0 = Conv2(17, 2*c)      # 3+3+3+3+1+4 = 17 channels

        self.down1 = Conv2(4*c, 4*c)
        self.down2 = Conv2(8*c, 8*c)
        self.down3 = Conv2(16*c, 16*c)
        # ... decoder layers

    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))
        # ... decoder forward pass

Each subsequent encoder stage concatenates its input with the corresponding context features from both frames (c0[i] and c1[i]), fusing image-level details with hierarchical semantic information.

Decoder Path and Skip Connections

The decoder employs transposed convolutions (deconv) to upsample features while progressively reducing channel counts. Skip connections concatenate encoder outputs (s0 through s3) with decoder features to recover spatial precision lost during downsampling:

        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)  # nn.Conv2d(c, 3, 3, 1, 1)

        return torch.sigmoid(x)

A final 3×3 convolution projects the feature map to 3 RGB channels, and a sigmoid activation constrains outputs to [0, 1]. In the complete pipeline, this output represents a residual correction that refines the initial warped-frame estimate.

Integration in the RIFE Pipeline

The top-level Model class in model/RIFE.py orchestrates the interaction between these modules:

# Extract flow-aware context for both frames

c0 = self.contextnet(img0, flow[:, :2])   # Context for frame 0

c1 = self.contextnet(img1, flow[:, 2:4])  # Context for frame 1

# Upsample flow to full resolution

flow = F.interpolate(flow, scale_factor=2.0, mode='bilinear', align_corners=False) * 2.0

# Fuse everything to produce the interpolated frame

refine_output = self.fusionnet(img0, img1, flow, c0, c1, flow_gt)

Here, self.contextnet instantiates either ContextNet or Contextnet depending on the model variant, while self.fusionnet is the Unet class. The UNet receives both raw images, their warped counterparts, the upsampled flow, and the four-scale context tensors to predict the final interpolation.

Practical Code Examples

Running ContextNet on a Single Frame

import torch
from model.refine import Contextnet  # or model.oldmodel.RIFE_HDv2.ContextNet

B, C, H, W = 1, 3, 256, 256
img = torch.randn(B, C, H, W).cuda()
flow = torch.randn(B, 2, H, W).cuda()  # Intermediate flow for this frame

ctx = Contextnet().cuda()
features = ctx(img, flow)  # List of 4 warped feature maps

for i, f in enumerate(features):
    print(f"Level {i}: {f.shape}")

# Output shapes: [1,16,128,128], [1,32,64,64], [1,64,32,32], [1,128,16,16]

Feeding UNet with Context and Warped Images

import torch
from model.refine import Unet, Contextnet, warp

B, C, H, W = 1, 3, 256, 256
img0 = torch.randn(B, C, H, W).cuda()
img1 = torch.randn(B, C, H, W).cuda()
flow = torch.randn(B, 4, H, W).cuda()  # Bidirectional flow [dx0, dy0, dx1, dy1]

# Extract context pyramids

ctx = Contextnet().cuda()
c0 = ctx(img0, flow[:, :2])
c1 = ctx(img1, flow[:, 2:4])

# Warp frames according to flow

warped0 = warp(img0, flow[:, :2])
warped1 = warp(img1, flow[:, 2:4])

# Initialize mask (learned during refinement)

mask = torch.zeros(B, 1, H, W).cuda()

# Run fusion network

unet = Unet().cuda()
output = unet(img0, img1, warped0, warped1, mask, flow, c0, c1)
print(output.shape)  # torch.Size([1, 3, 256, 256])

Summary

• ContextNet generates four-scale feature pyramids (1/4 to 1/32 resolution) by progressively downsampling input frames and warping features with intermediate optical flow.
• UNet fuses these context features with raw frames, warped frames, and bidirectional flow through an encoder-decoder architecture with skip connections.
• Both modules rely on the warp utility from model/warplayer.py for differentiable backward warping.
• The high-capacity ContextNet (c=32) resides in model/oldmodel/RIFE_HDv2.py, while the lightweight Contextnet (c=16) and Unet are defined in model/refine.py.
• Together, they enable the RIFE architecture to produce temporally consistent, high-fidelity intermediate frames by combining motion-aware feature alignment with deep feature fusion.

Frequently Asked Questions

What is the difference between ContextNet and Contextnet in the RIFE codebase?

ContextNet (capital N) refers to the high-capacity implementation in model/oldmodel/RIFE_HDv2.py that uses 32 base channels. Contextnet (lowercase n) is the lightweight variant in model/refine.py using 16 base channels, designed for faster inference with minimal quality degradation. Both perform identical multi-scale feature extraction and flow-guided warping operations.

Why does ContextNet warp features at every scale instead of only at full resolution?

Warping at each downsampling level ensures that context features remain spatially aligned with the intermediate frame despite large motion vectors. As resolution decreases, the flow magnitude is halved proportionally (* 0.5), allowing the network to capture both fine textures at high resolution and global motion patterns at lower resolutions. This hierarchical alignment prevents artifacts in scenes with complex motion.

How many input channels does the UNet encoder receive and what do they represent?

The first encoder layer (down0) receives 17 channels: 3 channels each for the two original RGB frames (img0, img1), 3 channels each for the flow-warped frames (warped_img0, warped_img1), 1 channel for the soft blending mask, and 4 channels for the bidirectional optical flow. Subsequent layers concatenate these with 2×, 4×, and 8× context features from both frames.

Can the ContextNet and UNet modules be used independently of the full RIFE model?

Yes, both modules are standard PyTorch nn.Module classes that can be instantiated and executed independently, as shown in the code examples above. However, they require properly formatted optical flow tensors (from IFNet or another estimator) to perform meaningful warping. The warp function from model/warplayer.py must be available in the environment for either module to function correctly.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}