File size: 5,482 Bytes
ebec118 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """LiquidInbetween: a small learned in-betweener for photos.
The architecture is RIFE's IFNet (Huang et al., ECCV 2022, MIT licence,
github.com/hzwer/ECCV2022-RIFE), extended with a timestep input so one pass
gives the frame at any t in (0, 1), not only the middle. The extra channel
starts at zero, so training begins exactly at RIFE's behaviour at t = 0.5.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
def conv(cin, cout, k=3, s=1, p=1):
return nn.Sequential(nn.Conv2d(cin, cout, k, s, p, bias=True), nn.PReLU(cout))
def base_grid(n, h, w, device, dtype):
gx = torch.linspace(-1.0, 1.0, w, device=device, dtype=dtype).view(1, 1, 1, w).expand(n, -1, h, -1)
gy = torch.linspace(-1.0, 1.0, h, device=device, dtype=dtype).view(1, 1, h, 1).expand(n, -1, -1, w)
return torch.cat([gx, gy], 1)
def warp(img, flow):
"""Backward warp: each output pixel reads `img` at its position + `flow` (pixels)."""
n, _, h, w = flow.shape
f = torch.cat([flow[:, 0:1] / ((w - 1.0) / 2.0), flow[:, 1:2] / ((h - 1.0) / 2.0)], 1)
# Clamping the grid gives border padding on every backend (MPS and Core ML too).
g = (base_grid(n, h, w, flow.device, flow.dtype) + f).clamp(-1.0, 1.0).permute(0, 2, 3, 1)
return F.grid_sample(img, g, mode="bilinear", padding_mode="zeros", align_corners=True)
class IFBlock(nn.Module):
def __init__(self, in_planes, c=90):
super().__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):
x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False)
flow = F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) / 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 = F.interpolate(self.conv1(feat), scale_factor=scale, mode="bilinear", align_corners=False) * scale
mask = F.interpolate(self.conv2(feat), scale_factor=scale, mode="bilinear", align_corners=False)
return flow, mask
class IFNet(nn.Module):
"""RIFE's IFNet. With `timestep`, every block also sees t (and the
mirrored pass 1 − t), so the frame at any t comes out in one pass."""
def __init__(self, timestep=False, c=90):
super().__init__()
self.timestep = timestep
extra = 1 if timestep else 0
self.block0 = IFBlock(7 + 4 + extra, c)
self.block1 = IFBlock(7 + 4 + extra, c)
self.block2 = IFBlock(7 + 4 + extra, c)
def forward(self, img0, img1, t=None, all_scales=False):
n, _, h, w = img0.shape
flow = torch.zeros(n, 4, h, w, device=img0.device, dtype=img0.dtype)
mask = torch.zeros(n, 1, h, w, device=img0.device, dtype=img0.dtype)
if self.timestep:
tt = (t if t is not None else torch.full((n, 1, 1, 1), 0.5, device=img0.device, dtype=img0.dtype)).expand(n, 1, h, w)
w0, w1 = img0, img1
outputs = []
for block, scale in ((self.block0, 4), (self.block1, 2), (self.block2, 1)):
a = [w0, w1, mask] + ([tt] if self.timestep else [])
b = [w1, w0, -mask] + ([1 - tt] if self.timestep else [])
f0, m0 = block(torch.cat(a, 1), flow, scale)
f1, m1 = block(torch.cat(b, 1), torch.cat((flow[:, 2:4], flow[:, :2]), 1), scale)
flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2
mask = mask + (m0 - m1) / 2
w0 = warp(img0, flow[:, :2])
w1 = warp(img1, flow[:, 2:4])
if all_scales:
m = torch.sigmoid(mask)
outputs.append(w0 * m + w1 * (1 - m))
m = torch.sigmoid(mask)
out = w0 * m + w1 * (1 - m)
return outputs if all_scales else out
def load_teacher(path):
"""RIFE's ECCV 2022 weights (safetensors) into IFNet()."""
from safetensors.torch import load_file
net = IFNet(timestep=False)
state = load_file(path)
state = {k.replace("module.", ""): v for k, v in state.items() if not k.startswith("block_tea")}
missing, unexpected = net.load_state_dict(state, strict=False)
assert not missing, missing
return net
def student_from_teacher(teacher):
"""IFNet(timestep=True) starting exactly as the teacher: the new t channel's
weights are zero."""
student = IFNet(timestep=True)
sd = student.state_dict()
for k, v in teacher.state_dict().items():
if k.endswith("conv0.0.0.weight"):
new = torch.zeros_like(sd[k])
# Inputs: warped0 (3), warped1 (3), mask (1) | t (1) | flow (4)
new[:, :7] = v[:, :7]
new[:, 8:] = v[:, 7:]
sd[k] = new
else:
sd[k] = v
student.load_state_dict(sd)
return student
|