ComfyUI/comfy_extras/rife_model/ifnet.py

import torch
import torch.nn as nn
import torch.nn.functional as F

import comfy.ops

ops = comfy.ops.disable_weight_init

_warp_grid_cache = {}
_WARP_GRID_CACHE_MAX = 4


def clear_warp_cache():
    _warp_grid_cache.clear()


def warp(img, flow):
    B, _, H, W = img.shape
    dtype = img.dtype
    img = img.float()
    flow = flow.float()
    cache_key = (H, W, flow.device)
    if cache_key not in _warp_grid_cache:
        if len(_warp_grid_cache) >= _WARP_GRID_CACHE_MAX:
            _warp_grid_cache.pop(next(iter(_warp_grid_cache)))
        grid_y, grid_x = torch.meshgrid(
            torch.linspace(-1.0, 1.0, H, device=flow.device, dtype=torch.float32),
            torch.linspace(-1.0, 1.0, W, device=flow.device, dtype=torch.float32),
            indexing="ij",
        )
        _warp_grid_cache[cache_key] = torch.stack((grid_x, grid_y), dim=0).unsqueeze(0)
    grid = _warp_grid_cache[cache_key].expand(B, -1, -1, -1)
    flow_norm = torch.cat([
        flow[:, 0:1] / ((W - 1) / 2),
        flow[:, 1:2] / ((H - 1) / 2),
    ], dim=1)
    grid = (grid + flow_norm).permute(0, 2, 3, 1)
    return F.grid_sample(img, grid, mode="bilinear", padding_mode="border", align_corners=True).to(dtype)


class Head(nn.Module):
    def __init__(self, out_ch=4, device=None, dtype=None, operations=ops):
        super().__init__()
        self.cnn0 = operations.Conv2d(3, 16, 3, 2, 1, device=device, dtype=dtype)
        self.cnn1 = operations.Conv2d(16, 16, 3, 1, 1, device=device, dtype=dtype)
        self.cnn2 = operations.Conv2d(16, 16, 3, 1, 1, device=device, dtype=dtype)
        self.cnn3 = operations.ConvTranspose2d(16, out_ch, 4, 2, 1, device=device, dtype=dtype)
        self.relu = nn.LeakyReLU(0.2, True)

    def forward(self, x):
        x = self.relu(self.cnn0(x))
        x = self.relu(self.cnn1(x))
        x = self.relu(self.cnn2(x))
        return self.cnn3(x)


class ResConv(nn.Module):
    def __init__(self, c, device=None, dtype=None, operations=ops):
        super().__init__()
        self.conv = operations.Conv2d(c, c, 3, 1, 1, device=device, dtype=dtype)
        self.beta = nn.Parameter(torch.ones((1, c, 1, 1), device=device, dtype=dtype))
        self.relu = nn.LeakyReLU(0.2, True)

    def forward(self, x):
        return self.relu(torch.addcmul(x, self.conv(x), self.beta))


class IFBlock(nn.Module):
    def __init__(self, in_planes, c=64, device=None, dtype=None, operations=ops):
        super().__init__()
        self.conv0 = nn.Sequential(
            nn.Sequential(
                operations.Conv2d(in_planes, c // 2, 3, 2, 1, device=device, dtype=dtype),
                nn.LeakyReLU(0.2, True),
            ),
            nn.Sequential(
                operations.Conv2d(c // 2, c, 3, 2, 1, device=device, dtype=dtype),
                nn.LeakyReLU(0.2, True),
            ),
        )
        self.convblock = nn.Sequential(
            *(ResConv(c, device=device, dtype=dtype, operations=operations) for _ in range(8))
        )
        self.lastconv = nn.Sequential(
            operations.ConvTranspose2d(c, 4 * 13, 4, 2, 1, device=device, dtype=dtype),
            nn.PixelShuffle(2),
        )

    def forward(self, x, flow=None, scale=1):
        x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear")
        if flow is not None:
            flow = F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear").div_(scale)
            x = torch.cat((x, flow), 1)
        feat = self.conv0(x)
        feat = self.convblock(feat)
        tmp = self.lastconv(feat)
        tmp = F.interpolate(tmp, scale_factor=scale, mode="bilinear")
        flow = tmp[:, :4] * scale
        mask = tmp[:, 4:5]
        feat = tmp[:, 5:]
        return flow, mask, feat


class IFNet(nn.Module):
    def __init__(self, head_ch=4, channels=None, device=None, dtype=None, operations=ops):
        super().__init__()
        if channels is None:
            channels = [192, 128, 96, 64, 32]
        self.encode = Head(out_ch=head_ch, device=device, dtype=dtype, operations=operations)
        block_in = [7 + 2 * head_ch] + [8 + 4 + 8 + 2 * head_ch] * 4
        self.blocks = nn.ModuleList([
            IFBlock(block_in[i], channels[i], device=device, dtype=dtype, operations=operations)
            for i in range(5)
        ])
        self.scale_list = [16, 8, 4, 2, 1]

    def get_dtype(self):
        return self.encode.cnn0.weight.dtype

    def forward(self, img0, img1, timestep=0.5):
        img0 = img0.clamp(0.0, 1.0)
        img1 = img1.clamp(0.0, 1.0)
        if not isinstance(timestep, torch.Tensor):
            timestep = torch.full((img0.shape[0], 1, img0.shape[2], img0.shape[3]),
                                  timestep, device=img0.device, dtype=img0.dtype)
        f0 = self.encode(img0)
        f1 = self.encode(img1)
        flow = mask = feat = None
        warped_img0 = img0
        warped_img1 = img1
        for i, block in enumerate(self.blocks):
            if flow is None:
                flow, mask, feat = block(
                    torch.cat((img0, img1, f0, f1, timestep), 1),
                    None, scale=self.scale_list[i],
                )
            else:
                wf0 = warp(f0, flow[:, :2])
                wf1 = warp(f1, flow[:, 2:4])
                fd, mask, feat = block(
                    torch.cat((warped_img0, warped_img1, wf0, wf1, timestep, mask, feat), 1),
                    flow, scale=self.scale_list[i],
                )
                flow = flow.add_(fd)
            warped_img0 = warp(img0, flow[:, :2])
            warped_img1 = warp(img1, flow[:, 2:4])
        mask = torch.sigmoid(mask)
        return torch.lerp(warped_img1, warped_img0, mask)


def detect_rife_config(state_dict):
    # Determine head output channels from encode.cnn3 (ConvTranspose2d)
    # ConvTranspose2d weight shape is (in_ch, out_ch, kH, kW)
    head_ch = state_dict["encode.cnn3.weight"].shape[1]

    # Read per-block channel widths from conv0 second layer output channels
    # conv0 is Sequential(Sequential(Conv2d, ReLU), Sequential(Conv2d, ReLU))
    # conv0.1.0.weight shape is (c, c//2, 3, 3)
    channels = []
    for i in range(5):
        key = f"blocks.{i}.conv0.1.0.weight"
        if key in state_dict:
            channels.append(state_dict[key].shape[0])

    if len(channels) != 5:
        raise ValueError(f"Unsupported RIFE model: expected 5 blocks, found {len(channels)}")

    return head_ch, channels